4_tasklens_img_classi.py

""" 
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)