[new] lens glasss, zmx io

Author: singer-yang

0_hello_deeplens.py

@@ -17,6 +17,7 @@
 def main():
     lens = GeoLens(filename='./lenses/camera/ef35mm_f2.0.json')
     # lens = GeoLens(filename='./lenses/cellphone/cellphone80deg.json')
+    # lens = GeoLens(filename='./lenses/zemax_double_gaussian.zmx')
     lens.analysis(render=True)
 
 if __name__=='__main__':

1_end2end_5lines.py

@@ -12,6 +12,7 @@
 import os
 import yaml
 import wandb
+import shutil
 import logging
 import random
 import string
@@ -63,9 +64,7 @@ def config():
     with open(f'{result_dir}/config.yml', 'w') as f:
         yaml.dump(args, f)
 
-    with open(f'{result_dir}/1_end2end_5lines.py', 'w') as f:
-        with open('1_end2end_5lines.py', 'r') as code:
-            f.write(code.read())
+    shutil.copy('1_end2end_5lines.py', f'{result_dir}/1_end2end_5lines.py')
 
     return args
 

2_autolens_rms.py

@@ -16,8 +16,6 @@
 import yaml
 import random
 import string
-
-
 from datetime import datetime
 from tqdm import tqdm
 from transformers import get_cosine_schedule_with_warmup
@@ -35,7 +33,7 @@ def config():
     characters = string.ascii_letters + string.digits
     random_string = ''.join(random.choice(characters) for i in range(4))
     current_time = datetime.now().strftime("%m%d-%H%M%S")
-    exp_name = current_time + '-Auto-Lens-Design-' + random_string
+    exp_name = current_time + '-AutoLens-RMS-' + random_string
     result_dir = f'./results/{exp_name}'
     os.makedirs(result_dir, exist_ok=True)
     args['result_dir'] = result_dir

4_tasklens_img_classi.py

@@ -0,0 +1,319 @@
+""" 
+Task-driven lens design for image classification. 
+
+We design a lens with from scratch with only image-classification loss. This makes sure no prior knowledge or classical lens design objective (spot size, PSF...) is used in the task-driven lens design. By doing this, we can explore "unseen" lens design space to find a lens that is optimal for a task, because we totally get rid of classical lens design!
+
+Technical Paper:
+    Xinge Yang, Yunfeng Nie, Fu Qiang and Wolfgang Heidrich, "Image Quality Is Not All You Want: Task-Driven Lens Design for Image Classification" Arxiv preprint 2023.
+
+This code and data is released under the Creative Commons Attribution-NonCommercial 4.0 International license (CC BY-NC.) In a nutshell:
+    # The license is only for non-commercial use (commercial licenses can be obtained from authors).
+    # The material is provided as-is, with no warranties whatsoever.
+    # If you publish any code, data, or scientific work based on this, please cite our work.
+"""
+import os
+import yaml
+import wandb
+import logging
+import random
+import string
+import timm
+import cv2 as cv
+from tqdm import tqdm
+from datetime import datetime
+from transformers import get_cosine_schedule_with_warmup
+
+import torch
+from torchvision.utils import save_image
+from torch.utils.data import DataLoader
+import torch.nn as nn
+import torchvision.transforms as transforms
+from torchvision.datasets import ImageFolder
+
+from deeplens import GeoLens
+from deeplens.utils import *
+from deeplens.optics.basics import *
+from deeplens.network.dataset import ImageDataset
+from deeplens.optics.render_psf import render_psf
+
+def config():
+    # ==> Config
+    with open('configs/4_tasklens.yml') as f:
+        args = yaml.load(f, Loader=yaml.FullLoader)
+
+    # ==> Result folder
+    characters = string.ascii_letters + string.digits
+    random_string = ''.join(random.choice(characters) for i in range(4))
+    result_dir = f'./results/' + datetime.now().strftime("%m%d-%H%M%S") + '-TaskLens' + '-' + random_string
+    args['result_dir'] = result_dir
+    os.makedirs(result_dir, exist_ok=True)
+    print(f'Result folder: {result_dir}')
+
+    if args['seed'] is None:
+        seed = random.randint(0, 100)
+        args['seed'] = seed
+    set_seed(args['seed'])
+    
+    # ==> Log
+    set_logger(result_dir)
+    if not args['DEBUG']:
+        raise Exception('Add your wandb logging config here.')
+
+    # ==> Device
+    num_gpus = torch.cuda.device_count()
+    args['num_gpus'] = num_gpus
+    device = torch.device(f"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)')
+
+    # ==> Save config
+    with open(f'{result_dir}/config.yml', 'w') as f:
+        yaml.dump(args, f)
+
+    with open(f'{result_dir}/4_tasklens_design.py', 'w') as f:
+        with open('4_tasklens_design.py', 'r') as code:
+            f.write(code.read())
+
+    return args
+
+
+def get_dataset(args):
+    dataset = args['train']['dataset']
+    img_res = args['train']['img_res']
+    bs = args['train']['bs']
+
+    # ==> Transforms
+    train_transform = transforms.Compose([
+        transforms.Resize(img_res),
+        transforms.RandomHorizontalFlip(),
+        transforms.TrivialAugmentWide(),
+        transforms.ToTensor(),
+        transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))
+    ])
+
+    val_transform = transforms.Compose([
+        transforms.Resize(img_res),
+        transforms.ToTensor(),
+        transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))
+    ])
+
+    # ==> Datset
+    if dataset == 'imagenet':
+        train_dataset = ImageFolder(root=args['imagenet_train_dir'], transform=train_transform)
+        val_dataset = ImageFolder(root=args['imagenet_val_dir'], transform=val_transform)
+    elif dataset == 'imagenet_local':
+        train_dataset = ImageFolder(root=args['imagenet_train_dir_local'], transform=train_transform)
+        val_dataset = ImageFolder(root=args['imagenet_val_dir_local'], transform=val_transform)
+    else:
+        raise NotImplementedError
+    
+    # ==> Data loader
+    train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=bs, shuffle=True)
+    val_loader = torch.utils.data.DataLoader(val_dataset, batch_size=bs, shuffle=False)
+    
+    return train_loader, val_loader
+
+
+def get_network(args):
+    if args['network']['model'] == 'resnet50':
+        net = timm.create_model(
+            'resnet50', 
+            pretrained=True, 
+            num_classes=1000
+        )
+    elif args['network']['model'] == 'swin_transformer':
+        net = timm.create_model(
+            'swin_base_patch4_window7_224_in22k', 
+            pretrained=True, 
+            num_classes=1000
+        )
+    elif args['network']['model'] == 'mobilenet':
+        net = timm.create_model(
+            'mobilenetv3_large_100', 
+            pretrained=True, 
+            num_classes=1000
+        )
+    elif args['network']['model'] == 'vit':
+        net = timm.create_model(
+            'vit_large_patch16_224_in21k', 
+            pretrained=True, 
+            num_classes=1000
+        )
+    else:
+        raise NotImplementedError
+
+    # Parallel
+    net = nn.DataParallel(net, device_ids=range(args['num_gpus']))
+    return net
+
+
+@torch.no_grad()
+def validate(lens, net, epoch, args, val_loader):
+    """ Test image classification accuracy.
+    """
+    # Parameters
+    device = args['device']
+    result_dir = args['result_dir']
+    depth = args['train']['depth']
+    bs = args['train']['bs']
+    ks = args['train']['psf_ks']
+    psf_grid = args['train']['psf_grid']
+    points = lens.point_source_grid(depth=depth, grid=psf_grid*2-1, quater=True).reshape(-1, 3)
+
+    # Scores
+    correct = 0.0
+    total = 0.0
+
+    # Calculate PSFs
+    psf = lens.psf_rgb(points=points, ks=ks, spp=4096)
+
+    # Loop over the validation set in batches
+    for _, (img_org, labels) in tqdm(enumerate(val_loader)):
+        if img_org.shape[0] != bs:
+            continue
+
+        # Get images and labels
+        img_org = img_org.to(device)
+        labels = labels.to(device)
+
+        # Render image with PSF map
+        img_render = render_psf(img_org, psf)
+        img_render = torch.cat(img_render)
+        labels = labels.repeat(psf_grid**2)
+
+        # Forward pass and prediction
+        outputs = net(img_render)
+        _, predicted = torch.max(outputs.data, 1)
+
+        # Update accuracy statistics
+        total += labels.size(0)
+        correct += (predicted == labels).sum().item()
+
+    # Print validation accuracy
+    acc = correct / total
+    if acc > args['val_acc']:
+        args['val_acc'] = acc
+        logging.info(f'Best epoch is {epoch}, best Val acc is {acc}.')
+        torch.save(net.state_dict(), f'{result_dir}/classi_model_best.pth')
+
+    logging.info('Validation Accuracy: {:.2f}%'.format(100 * acc))
+    if not args['DEBUG']:
+        wandb.log({"classi_acc":acc})
+
+
+def train(args, lens, net):
+    device = args['device']
+    result_dir = args['result_dir']
+    bs = args['train']['bs']
+    ks = args['train']['psf_ks']
+    psf_grid = args['train']['psf_grid']
+    spp = args['train']['spp']
+    depth = args['train']['depth']
+    lens_lrs = [float(i) for i in args['lens']['lr']]
+    args['val_acc'] = 0
+
+    # ==> Dataset
+    train_loader, val_loader = get_dataset(args)
+    batchs = len(train_loader)
+    epochs = args['train']['epochs']
+    
+    # ==> Optimizer and scheduler
+    lens_optim = lens.get_optimizer(lr=lens_lrs)
+    lens_sche = get_cosine_schedule_with_warmup(lens_optim, num_warmup_steps=500, num_training_steps=batchs*epochs)
+    # # Uncomment for End-to-End lens-network co-design
+    # net_optim = torch.optim.Adam(net.parameters(), lr=1e-4)
+    # net_sche = get_cosine_schedule_with_warmup(net_optim, num_warmup_steps=500, num_training_steps=batchs*epochs)
+
+    # ==> Loss
+    cri_classi = nn.CrossEntropyLoss()
+
+    # ==> Training
+    logging.info(f'==> Start training.')
+    points = lens.point_source_grid(depth=depth, grid=psf_grid, quater=True).reshape(-1, 3)
+    for epoch in range(args['train']['epochs'] + 1):
+        
+        # =============================
+        # Evaluation 
+        # =============================
+        if epoch % 1 == 0 and epoch > 0:
+            net.eval()
+            lens.correct_shape()
+            lens.write_lens_json(f'{result_dir}/epoch{epoch}.json')
+            lens.analysis(f'{result_dir}/epoch{epoch}', render=False)
+            validate(lens, net, epoch, args, val_loader)
+
+        # =============================
+        # Training
+        # =============================
+        net.train()
+
+        # ==> Task-driven lens design: a well-trained network serves as lens design objective
+        for ii, (img_org, labels) in tqdm(enumerate(train_loader)):
+
+            # Continue is wrong batch size
+            if img_org.shape[0] != bs:
+                continue
+            
+            # Get images and labels
+            img_org = img_org.to(device)
+            labels = labels.to(device)
+
+            # Option 1: Render image with PSF map
+            psf = lens.psf_rgb(points=points, ks=ks, center=False, spp=spp)  # [N, 3, ks, ks]
+            img_render = []
+            for psf_idx in range(psf.shape[0]):
+                img_render.append(render_psf(img_org, psf[psf_idx, ...]))
+            img_render = torch.cat(img_render) # [N * B, 3, sensor_res, sensor_res]
+            labels = labels.repeat(psf.shape[0])
+
+            # Option 2: Render image with ray tracing
+            # img_render = lens.render(img_org)
+
+            # Image classification
+            labels_pred = net(img_render)
+
+            # Loss
+            L_classi = cri_classi(labels_pred, labels)
+            L_reg = lens.loss_self_intersec() #+ lens.loss_ray_angle()
+
+            L = L_classi + 0.02 * L_reg
+
+            # Update
+            lens_optim.zero_grad()
+            # net_optim.zero_grad()
+            L.backward()
+            lens_optim.step()
+            # net_optim.step()
+            lens_sche.step()
+            # net_sche.step()
+
+            if not args['DEBUG']:
+                wandb.log({"loss_class": L_classi.detach().item()})
+
+            # Print statistics every 1000 batches
+            if ii % 100 == 0 and ii > 0:
+                logging.info('Epoch [{}/{}], Batch [{}/{}], Loss: {:.4f}'.format(epoch+1, args['train']['epochs'], ii, len(train_loader), L.item()))
+                lens.correct_shape()
+                lens.write_lens_json(f'{result_dir}/epoch{epoch}_batch{ii}.json')
+                lens.analysis(f'{result_dir}/epoch{epoch}_batch{ii}', render=False)
+        
+        logging.info(f'Epoch{epoch+1} finishs.')
+
+
+if __name__=='__main__':
+    args = config()
+    
+    # Lens
+    lens = GeoLens(filename=args['lens']['path'], sensor_res=args['lens']['sensor_res'], device=args['device'])
+    lens.set_target_fov_fnum(hfov=args['lens']['target_hfov'], fnum=args['lens']['target_fnum'])
+    lens.write_lens_json(f'{args["result_dir"]}/epoch0.json')
+    lens.analysis(f'{args["result_dir"]}/epoch0', render=False, zmx_format=True)
+
+    # Network
+    net = get_network(args)
+    for param in net.parameters():
+        param.requires_grad = False
+    net = net.to(args['device'])
+
+    # End-to-end lens-network co-design
+    train(args, lens, net)