Skip to content

Anamorphics #25

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
wants to merge 1 commit into
base: main
Choose a base branch
from
Draft

Anamorphics #25

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
34 changes: 32 additions & 2 deletions 0_hello_deeplens.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,10 +17,40 @@


def main():
lens = GeoLens(filename="./lenses/camera/ef35mm_f2.0.json")
lens = GeoLens(filename="./lenses/camera/anamorphic_50mm.json")
# lens = GeoLens(filename='./lenses/cellphone/cellphone80deg.json')
# lens = GeoLens(filename='./lenses/zemax_double_gaussian.zmx')
lens.analysis(render=True)
lens.analysis(
f"./initial_lens",
zmx_format=True,
multi_plot=False,
)
lens.optimize(
lrs=[6e-4, 1e-4, 0.1, 1e-4],
decay=0.02,
iterations=5000,
centroid=True,
importance_sampling=True,
optim_mat=True,
shape_control=True,
anamorphic=True,
test_per_iter=20,
result_dir="./result",
)

# =====> 3. Analyze final result
lens.prune_surf(expand_surf=0.02)
Copy link
Owner

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It is better to expand more for large size lenses, for example, expand_surf = 0.1 or even 0.2

lens.post_computation()

logging.info(
f"Actual: diagonal FOV {lens.hfov}, r sensor {lens.r_sensor}, F/{lens.fnum}."
)
lens.write_lens_json(f"{result_dir}/final_lens.json")
lens.analysis(save_name=f"{result_dir}/final_lens", zmx_format=True)

# =====> 4. Create video
create_video_from_images(f"{result_dir}", f"{result_dir}/autolens.mp4", fps=10)



if __name__ == "__main__":
Expand Down
13 changes: 10 additions & 3 deletions 1_end2end_5lines.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -66,11 +66,11 @@ def config():
raise Exception("Add your wandb logging config here.")

# ==> Device
num_gpus = torch.cuda.device_count()
num_gpus = 1
args["num_gpus"] = num_gpus
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
device = torch.device("cuda" if torch.cuda.is_available() else "mps")
args["device"] = device
logging.info(f"Using {num_gpus} {torch.cuda.get_device_name(0)} GPU(s)")
#logging.info(f"Using {num_gpus} {torch.cuda.get_device_name(0)} GPU(s)")

# ==> Save config and original code
with open(f"{result_dir}/config.yml", "w") as f:
Expand Down Expand Up @@ -100,6 +100,7 @@ def end2end_train(lens, net, args):

train_set = ImageDataset(args["train"]["train_dir"], lens.sensor_res)
train_loader = DataLoader(train_set, batch_size=args["train"]["bs"])
logging.info(f'train_loader: {train_loader}')

# ==> Network optimizer
batchs = len(train_loader)
Expand Down Expand Up @@ -134,18 +135,22 @@ def end2end_train(lens, net, args):
# ==> Train 1 epoch
for img_org in tqdm(train_loader):
img_org = img_org.to(device)
logging.info(f'img_org: {img_org}')

# => Render image
# ========================================
# Line 3: plug-and-play diff-rendering
# ========================================
img_render = lens.render(img_org)
logging.info(f'img_render: {img_render}')

# => Image restoration
img_rec = net(img_render)
logging.info(f'img_rec: {img_rec}')

# => Loss
L_rec = cri_l1(img_rec, img_org)
logging.info(f'L_rec: {L_rec}')

# => Back-propagation
net_optim.zero_grad()
Expand Down Expand Up @@ -225,9 +230,11 @@ def end2end_train(lens, net, args):
# Line 1: load a lens
# ========================================
lens = GeoLens(filename=args["lens"]["path"])
logging.info(f'lens: {lens}')
lens.change_sensor_res(args["train"]["img_res"])
net = UNet()
net = net.to(lens.device)
logging.info(f'net: {net}')
if args["network"]["pretrained"]:
net.load_state_dict(torch.load(args["network"]["pretrained"]))

Expand Down
43 changes: 37 additions & 6 deletions 2_autolens_rms.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -58,11 +58,11 @@ def config():
logging.info(f"EXP: {args['EXP_NAME']}")

# Device
num_gpus = torch.cuda.device_count()
num_gpus = 1 #torch.cuda.device_count()
args["num_gpus"] = num_gpus
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
device = torch.device("mps" if torch.backends.mps.is_available() else "cpu")
args["device"] = device
logging.info(f"Using {num_gpus} {torch.cuda.get_device_name(0)} GPU(s)")
#logging.info(f"Using {num_gpus} {torch.cuda.get_device_name(0)} GPU(s)")

# ==> Save config and original code
with open(f"{result_dir}/config.yml", "w") as f:
Expand Down Expand Up @@ -153,9 +153,13 @@ def curriculum_design(
center_p = center_p.unsqueeze(-2).repeat(1, 1, spp, 1)

# =======================================
# Optimize lens by minimizing rms
# Optimize lens by minimizing rms and anamorphic squeeze
# =======================================
loss_rms = []
loss_anamorphic_list = []
delta = 1 # object-space offset (ensure consistent units)
target_squeeze = 1.5 # desired horizontal/vertical magnification ratio

for j, wv in enumerate(WAVE_RGB):
# Ray tracing to sensor
ray = rays_backup[j].clone()
Expand Down Expand Up @@ -185,13 +189,40 @@ def curriculum_design(
)
loss_rms.append(l_rms)

# ---- Anamorphic (Squeeze) Loss Computation ----
# Horizontal offset ray.
ray_offset_h = self.get_cached_rays(depth=depth, wvln=wv, offset='h')
ray_offset_h = self.trace2sensor(ray_offset_h)
xy_offset_h = ray_offset_h.o[..., :2]
ra_xy_offset_h = ray_offset_h.ra.clone().detach()
xy_offset_h_norm = (xy_offset_h - center_p) * ra_xy_offset_h.unsqueeze(-1)

# Vertical offset ray.
ray_offset_v = self.get_cached_rays(depth=depth, wvln=wv, offset='v')
ray_offset_v = self.trace2sensor(ray_offset_v)
xy_offset_v = ray_offset_v.o[..., :2]
ra_xy_offset_v = ray_offset_v.ra.clone().detach()
xy_offset_v_norm = (xy_offset_v - center_p) * ra_xy_offset_v.unsqueeze(-1)

# Compute effective magnifications.
M_x = torch.mean(torch.abs(xy_offset_h_norm[..., 0])) / delta
M_y = torch.mean(torch.abs(xy_offset_v_norm[..., 1])) / delta
print(f'Magnifications: {M_x}, {M_y}')
ratio = M_x / (M_y + EPSILON)
loss_anamorphic_list.append((ratio - target_squeeze) ** 2)

# RMS loss for all wavelengths
loss_rms = sum(loss_rms) / len(loss_rms)
#loss_anamorphic = torch.mean(torch.stack(loss_anamorphic_list))
loss_anamorphic = sum(loss_anamorphic_list) / len(loss_anamorphic_list)

# Lens design constraint
loss_reg = self.loss_reg()
w_reg = 0.1
L_total = loss_rms + w_reg * loss_reg
# Adding for anamorphics. 1.5x squeeze hardcoded for now
w_anamorphic = 1.0
print(f"Losses: {loss_rms}, {w_reg * loss_reg}, {w_anamorphic * loss_anamorphic}")
L_total = loss_rms + w_reg * loss_reg + w_anamorphic * loss_anamorphic

# Gradient-based optimization
optimizer.zero_grad()
Expand Down Expand Up @@ -236,7 +267,7 @@ def curriculum_design(
lrs=[float(lr) for lr in args["lrs"]],
decay=float(args["decay"]),
iterations=5000,
test_per_iter=50,
test_per_iter=3,
optim_mat=True,
match_mat=False,
shape_control=True,
Expand Down
6 changes: 3 additions & 3 deletions 6_hybridlens_design.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -55,11 +55,11 @@ def config():
raise Exception("Add your wandb logging config here.")

# ==> Device
num_gpus = torch.cuda.device_count()
num_gpus = torch.cuda.device_count() if torch.cuda.is_available() else 0
args["num_gpus"] = num_gpus
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
args["device"] = device
logging.info(f"Using {num_gpus} {torch.cuda.get_device_name(0)} GPU(s)")
# logging.info(f"Using {num_gpus} {torch.cuda.get_device_name(0)} GPU(s)")

# ==> Save config
with open(f"{result_dir}/config.yml", "w") as f:
Expand All @@ -76,7 +76,7 @@ def main(args):
# Create a hybrid refractive-diffractive lens
lens = HybridLens(lens_path="./lenses/hybridlens/a489_doe.json")
lens.refocus(foc_dist=-1000.0)
lens.double()
#lens.double()

# PSF optimization loop to focus blue light
optimizer = lens.get_optimizer(doe_lr=0.1, lens_lr=[1e-4, 1e-4, 1e-1, 1e-5])
Expand Down
16 changes: 8 additions & 8 deletions configs/2_auto_lens_design.yml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,16 +2,16 @@ DEBUG: True
seed: ~

# experiment settings
EXP_NAME: 'Auto lens design'
EXP_NAME: 'Auto lens design - anamorphic'

# lens target example 1 (camera lens)
foclen: 85.0
fov: 40.0
foclen: 50.0
fov: 48.0
fnum: 4.0
flange: 18.0
thickness: 120.0
lens_type: [["Spheric", "Spheric"], ["Spheric", "Spheric"], ["Spheric", "Spheric", "Spheric"], ["Aperture"], ["Spheric", "Spheric"], ["Spheric", "Aspheric"], ["Spheric", "Aspheric"]]
lrs: [5e-4, 1e-3, 1e-1, 1e-3]
flange: 27.75
thickness: 111.906
lens_type: [["Anamorphic", "Anamorphic", "Anamorphic"], ["Spheric", "Spheric", "Spheric"], ["Spheric", "Spheric"], ["Aperture"], ["Spheric", "Spheric"], ["Spheric", "Spheric", "Spheric"], ["Anamorphic", "Anamorphic"], ["Anamorphic", "Anamorphic"]]
lrs: [2e-3, 2e-3, 1e-1, 1e-3]
decay: 0.001

# # lens target example 2 (mobile lens)
Expand All @@ -22,4 +22,4 @@ decay: 0.001
# thickness: 9.0
# lens_type: [["Aperture"], ["Aspheric", "Aspheric"], ["Aspheric", "Aspheric"], ["Aspheric", "Aspheric"], ["Aspheric", "Aspheric"], ["Aspheric", "Aspheric"]]
# lrs: [3e-4, 1e-4, 1e-1, 1e-2]
# decay: 0.01
# decay: 0.01
25 changes: 25 additions & 0 deletions configs/2_auto_lens_design_original.yml
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,25 @@
DEBUG: True
seed: ~

# experiment settings
EXP_NAME: 'Auto lens design'

# lens target example 1 (camera lens)
foclen: 85.0
fov: 40.0
fnum: 4.0
flange: 18.0
thickness: 120.0
lens_type: [["Spheric", "Spheric"], ["Spheric", "Spheric"], ["Spheric", "Spheric", "Spheric"], ["Aperture"], ["Spheric", "Spheric"], ["Spheric", "Aspheric"], ["Spheric", "Aspheric"]]
lrs: [5e-4, 1e-3, 1e-1, 1e-3]
decay: 0.001

# # lens target example 2 (mobile lens)
# foclen: 6.0
# fov: 70.0
# fnum: 2.0
# flange: 1.0
# thickness: 9.0
# lens_type: [["Aperture"], ["Aspheric", "Aspheric"], ["Aspheric", "Aspheric"], ["Aspheric", "Aspheric"], ["Aspheric", "Aspheric"], ["Aspheric", "Aspheric"]]
# lrs: [3e-4, 1e-4, 1e-1, 1e-2]
# decay: 0.01
Loading