initail upload

README.md

@@ -0,0 +1,185 @@
+# CAM-Seg: A Continuous-valued Embedding Approach for Semantic Image Generation
+
+**Official PyTorch Implementation**
+
+This is a PyTorch/GPU implementation of the paper [CAM-Seg: A Continuous-valued Embedding Approach for Semantic Image Generation](https://arxiv.org/abs/2503.15617)
+
+```
+@article{ahmed2025cam,
+  title={CAM-Seg: A Continuous-valued Embedding Approach for Semantic Image Generation},
+  author={Ahmed, Masud and Hasan, Zahid and Haque, Syed Arefinul and Faridee, Abu Zaher Md and Purushotham, Sanjay and You, Suya and Roy, Nirmalya},
+  journal={arXiv preprint arXiv:2503.15617},
+  year={2025}
+}
+```
+
+## Abstract
+Traditional transformer-based semantic segmentation relies on quantized embeddings. However, our analysis reveals that autoencoder accuracy on segmentation mask using quantized embeddings (e.g. VQ-VAE) is 8\% lower than continuous-valued embeddings  (e.g. KL-VAE). Motivated by this, we propose a continuous-valued embedding framework for semantic segmentation. By reformulating semantic mask generation as a continuous image-to-embedding diffusion process, our approach eliminates the need for discrete latent representations while preserving fine-grained spatial and semantic details. Our key contribution includes a diffusion-guided autoregressive transformer that learns a continuous semantic embedding space by modeling long-range dependencies in image features. Our framework contains a unified architecture combining a VAE encoder for continuous feature extraction, a diffusion-guided transformer for conditioned embedding generation, and a VAE decoder for semantic mask reconstruction. Our setting facilitates zero-shot domain adaptation capabilities enabled by the continuity of the embedding space. Experiments across diverse datasets (e.g., Cityscapes and domain-shifted variants) demonstrate state-of-the-art robustness to distribution shifts, including adverse weather (e.g., fog, snow) and viewpoint variations. Our model also exhibits strong noise resilience, achieving robust performance ($\approx$ 95\% AP compared to baseline) under gaussian noise, moderate motion blur, and moderate brightness/contrast variations, while experiencing only a moderate impact ($\approx$ 90\% AP compared to baseline) from 50\% salt and pepper noise, saturation and hue shifts.
+
+## Result
+Trained on Cityscape dataset and tested on SemanticKITTI, ACDC, CADEdgeTune dataset
+<p align="center">
+  <img src="demo/qualitative.png" width="720">
+</p>
+
+Quantitative results of semantic segmentation under various noise conditions
+<p align="center">
+    <table>
+      <tr>
+        <td align="center"><img src="demo/saltpepper_noise.png" width="200"/><br>Salt & Pepper Noise</td>
+        <td align="center"><img src="demo/motion_blur.png" width="200"/><br>Motion Blur</td>
+        <td align="center"><img src="demo/gaussian_noise.png" width="200"/><br>Gaussian Noise</td>
+        <td align="center"><img src="demo/gaussian_blur.png" width="200"/><br>Gaussian Blur</td>
+      </tr>
+      <tr>
+        <td align="center"><img src="demo/brightness.png" width="200"/><br>Brightness Variation</td>
+        <td align="center"><img src="demo/contrast.png" width="200"/><br>Contrast Variation</td>
+        <td align="center"><img src="demo/saturation.png" width="200"/><br>Saturation Variation</td>
+        <td align="center"><img src="demo/hue.png" width="200"/><br>Hue Variation</td>
+      </tr>
+    </table>
+</p>
+
+## Prerequisite
+To install the docker environment, first edit the `docker_env/Makefile`:
+```
+IMAGE=img_name/dl-aio
+CONTAINER=containter_name
+AVAILABLE_GPUS='0,1,2,3'
+LOCAL_JUPYTER_PORT=18888
+LOCAL_TENSORBOARD_PORT=18006
+PASSWORD=yourpassword
+WORKSPACE=workspace_directory
+```
+- Edit the `img_name`, `containter_name`, `available_gpus`, `jupyter_port`, `tensorboard_port`, `password`, `workspace_directory`
+
+1. For the first time run the following commands in terminal:
+```
+cd docker_env
+make docker-build
+make docker-run
+```
+2. or further use to docker environment
+- To stop the environmnet: `make docker-stop`
+- To resume the environmente: `make docker-resume`
+
+For coding open a web browser `ip_address:jupyter_port` e.g.,`http://localhost:18888`
+
+## Dataset
+Four Dataset is used in the work
+1. [Cityscapes Dataset](https://www.cityscapes-dataset.com/)
+2. [KITTI Dataset](https://www.cvlibs.net/datasets/kitti/eval_step.php)
+3. [ACDC Dataset](https://acdc.vision.ee.ethz.ch/)
+4. [CAD-EdgeTune Dataset](https://ieee-dataport.org/documents/cad-edgetune)
+
+**Modify the trainlist and vallist file to edit train and test split**
+
+### Dataset structure
+- Cityscapes Dataset
+```
+|-CityScapes
+|----leftImg8bit #contians the RGB images
+|----gtFine #contains semantic segmentation labels
+|----trainlist.txt #image list used for training
+|----vallist.txt #image list used for testing
+|----cityscape.yaml #configuration file for Cityscapes dataset
+```
+
+- ACDC Dataset
+```
+|-ACDC
+|----rgb_anon #contians the RGB images
+|----gt #contains semantic segmentation labels
+|----vallist_fog.txt #image list used for testing fog data
+|----vallist_rain.txt #image list used for testing rain data
+|----vallist_snow.txt #image list used for testing snow data
+|----acdc.yaml #configuration file for ACDC dataset
+```
+
+## Weights
+To download the pretrained weights please visit [Hugging Face Repo](https://huggingface.co/mahmed10/CAM-Seg)
+- **LDM model** Pretrained model from Rombach et al.'s Latent Diffusion Models is used [Link](https://huggingface.co/mahmed10/CAM-Seg/resolve/main/pretrained_models/vae/modelf16.ckpt)
+- **MAR model** Following mar model is used
+
+|Training Data|Model|Params|Link|
+|-------------|-----|------|----|
+|Cityscapes | Mar-base| 217M|[link](https://huggingface.co/mahmed10/CAM-Seg/resolve/main/pretrained_models/mar/city768.16.pth)|
+
+
+Download this weight files and organize as follow
+```
+|-pretrained_models
+|----mar
+|--------city768.16.pth
+|----vae
+|--------modelf16.ckpt
+```
+
+**Alternative code to automatically download pretrain weights**
+```
+import os
+import requests
+
+# Define URLs and file paths
+files_to_download = {
+    "https://huggingface.co/mahmed10/CAM-Seg/resolve/main/pretrained_models/vae/modelf16.ckpt":
+        "pretrained_models/vae/modelf16.ckpt",
+    "https://huggingface.co/mahmed10/CAM-Seg/resolve/main/pretrained_models/mar/city768.16.pth":
+        "pretrained_models/mar/city768.16.pth"
+}
+
+for url, path in files_to_download.items():
+    os.makedirs(os.path.dirname(path), exist_ok=True)
+
+    print(f"Downloading from {url}...")
+    response = requests.get(url, stream=True)
+    if response.status_code == 200:
+        with open(path, 'wb') as f:
+            for chunk in response.iter_content(chunk_size=8192):
+                f.write(chunk)
+        print(f"Saved to {path}")
+    else:
+        print(f"Failed to download from {url}, status code {response.status_code}")
+```
+
+## Validation
+Open the `validation.ipnyb` file
+
+Edit the **Block 6** to select which dataset is to use for validation
+
+```
+dataset_train = cityscapes.CityScapes('dataset/CityScapes/vallist.txt', data_set= 'val', transform=transform_train,seed=36, img_size=768)
+# dataset_train = umbc.UMBC('dataset/UMBC/all.txt', data_set= 'val', transform=transform_train,seed=36, img_size=768)
+# dataset_train = acdc.ACDC('dataset/ACDC/vallist_fog.txt', data_set= 'val', transform=transform_train,seed=36, img_size=768)
+# dataset_train = semantickitti.SemanticKITTI('dataset/SemanticKitti/vallist.txt', data_set= 'val', transform=transform_train, seed=36, img_size=768)
+```
+
+Run all the blocks
+
+## Training
+
+### From Scratch
+
+Run the following code in terminal
+```
+torchrun --nproc_per_node=4 train.py
+```
+
+it will save checkpoint in `output_dir/year.month.day.hour.min` folder, for e.g. `output_dir/2025.05.09.02.27`
+
+### Resume Training
+
+Run the following code in terminal
+```
+torchrun --nproc_per_node=4 train.py --resume year.month.day.hour.min
+```
+
+Here is an example code
+```
+torchrun --nproc_per_node=4 train.py --resume 2025.05.09.02.27
+```
+
+## Acknowlegement
+The code is developed on top following codework
+1. [latent-diffusion](https://github.com/CompVis/latent-diffusion)
+2. [mar](https://github.com/LTH14/mar)

get_parser.py

@@ -0,0 +1,121 @@
+import argparse
+from pathlib import Path
+import yaml
+def get_args_parser():
+    parser = argparse.ArgumentParser('MAR training with Diffusion Loss', add_help=False)
+    parser.add_argument('--batch_size', default=16, type=int,
+                        help='Batch size per GPU (effective batch size is batch_size * # gpus')
+    parser.add_argument('--epochs', default=2000, type=int)
+
+    # Model parameters
+    parser.add_argument('--model', default='mar_base', type=str, metavar='MODEL',
+                        help='Name of model to train')
+    parser.add_argument('--ckpt_path', default="pretrained_models/mar/city768.16.pth", type=str,
+                        help='model checkpoint path')
+
+    # VAE parameters
+    parser.add_argument('--img_size', default=768, type=int,
+                        help='images input size')
+    parser.add_argument('--vae_path', default="pretrained_models/vae/modelf16.ckpt", type=str,
+                        help='images input size')
+    parser.add_argument('--vae_embed_dim', default=16, type=int,
+                        help='vae output embedding dimension')
+    parser.add_argument('--vae_stride', default=16, type=int,
+                        help='tokenizer stride, default use KL16')
+    parser.add_argument('--patch_size', default=1, type=int,
+                        help='number of tokens to group as a patch.')
+    parser.add_argument('--config', default="ldm/config.yaml", type=str,
+                        help='vae model configuration file')
+
+    # Generation parameters
+    parser.add_argument('--num_iter', default=64, type=int,
+                        help='number of autoregressive iterations to generate an image')
+    parser.add_argument('--num_images', default=3000, type=int,
+                        help='number of images to generate')
+    parser.add_argument('--cfg', default=1.0, type=float, help="classifier-free guidance")
+    parser.add_argument('--cfg_schedule', default="linear", type=str)
+    parser.add_argument('--label_drop_prob', default=0.1, type=float)
+    parser.add_argument('--eval_freq', type=int, default=40, help='evaluation frequency')
+    parser.add_argument('--save_last_freq', type=int, default=5, help='save last frequency')
+    parser.add_argument('--online_eval', action='store_true')
+    parser.add_argument('--evaluate', action='store_true')
+    parser.add_argument('--eval_bsz', type=int, default=64, help='generation batch size')
+
+    # Optimizer parameters
+    parser.add_argument('--weight_decay', type=float, default=0.02,
+                        help='weight decay (default: 0.02)')
+
+    parser.add_argument('--grad_checkpointing', action='store_true')
+    parser.add_argument('--lr', type=float, default=None, metavar='LR',
+                        help='learning rate (absolute lr)')
+    parser.add_argument('--blr', type=float, default=1e-4, metavar='LR',
+                        help='base learning rate: absolute_lr = base_lr * total_batch_size / 256')
+    parser.add_argument('--min_lr', type=float, default=0., metavar='LR',
+                        help='lower lr bound for cyclic schedulers that hit 0')
+    parser.add_argument('--lr_schedule', type=str, default='constant',
+                        help='learning rate schedule')
+    parser.add_argument('--warmup_epochs', type=int, default=100, metavar='N',
+                        help='epochs to warmup LR')
+    parser.add_argument('--ema_rate', default=0.9999, type=float)
+
+    # MAR params
+    parser.add_argument('--mask_ratio_min', type=float, default=0.7,
+                        help='Minimum mask ratio')
+    parser.add_argument('--grad_clip', type=float, default=3.0,
+                        help='Gradient clip')
+    parser.add_argument('--attn_dropout', type=float, default=0.1,
+                        help='attention dropout')
+    parser.add_argument('--proj_dropout', type=float, default=0.1,
+                        help='projection dropout')
+    parser.add_argument('--buffer_size', type=int, default=64)
+
+    # Diffusion Loss params
+    parser.add_argument('--diffloss_d', type=int, default=6)
+    parser.add_argument('--diffloss_w', type=int, default=1024)
+    parser.add_argument('--num_sampling_steps', type=str, default="100")
+    parser.add_argument('--diffusion_batch_mul', type=int, default=4)
+    parser.add_argument('--temperature', default=1.0, type=float, help='diffusion loss sampling temperature')
+
+    # Dataset parameters
+    parser.add_argument('--output_dir', default='./output_dir',
+                        help='path where to save, empty for no saving')
+    parser.add_argument('--log_dir', default='./output_dir',
+                        help='path where to tensorboard log')
+    parser.add_argument('--device', default='cuda',
+                        help='device to use for training / testing')
+    parser.add_argument('--seed', default=1, type=int)
+    parser.add_argument('--resume', default=None,#'pretrained_models/mar/mar_base',
+                        help='resume from checkpoint')
+
+    parser.add_argument('--start_epoch', default=0, type=int, metavar='N',
+                        help='start epoch')
+    parser.add_argument('--num_workers', default=10, type=int)
+    parser.add_argument('--pin_mem', action='store_true',
+                        help='Pin CPU memory in DataLoader for more efficient (sometimes) transfer to GPU.')
+    parser.add_argument('--no_pin_mem', action='store_false', dest='pin_mem')
+    parser.set_defaults(pin_mem=True)
+
+    # distributed training parameters
+    parser.add_argument('--world_size', default=1, type=int,
+                        help='number of distributed processes')
+    parser.add_argument('--local_rank', default=-1, type=int)
+    parser.add_argument('--dist_on_itp', action='store_true')
+    parser.add_argument('--dist_url', default='env://',
+                        help='url used to set up distributed training')
+
+    # caching latents
+    parser.add_argument('--use_cached', action='store_true', dest='use_cached',
+                        help='Use cached latents')
+    parser.set_defaults(use_cached=False)
+    parser.add_argument('--cached_path', default='', help='path to cached latents')
+
+    return parser
+
+args = get_args_parser()
+args = args.parse_args()
+Path(args.output_dir).mkdir(parents=True, exist_ok=True)
+args.log_dir = args.output_dir
+
+with open(args.config, "r") as f:
+    config = yaml.safe_load(f)
+args.ddconfig = config["ddconfig"]

train.py

@@ -0,0 +1,381 @@
+import argparse
+import datetime
+import numpy as np
+import os
+import time
+from pathlib import Path
+import yaml
+import glob
+
+import torch
+import torch.backends.cudnn as cudnn
+from torch.utils.tensorboard import SummaryWriter
+import torchvision.transforms as transforms
+import torchvision.datasets as datasets
+from data import cityscapes
+
+from util.crop import center_crop_arr
+import util.misc as misc
+from util.misc import NativeScalerWithGradNormCount as NativeScaler
+from util.loader import CachedFolder
+
+from models.vae import AutoencoderKL
+from models import mar
+import copy
+from tqdm import tqdm
+
+import util.lr_sched as lr_sched
+
+import logging
+
+
+
+def update_ema(target_params, source_params, rate=0.99):
+    """
+    Update target parameters to be closer to those of source parameters using
+    an exponential moving average.
+
+    :param target_params: the target parameter sequence.
+    :param source_params: the source parameter sequence.
+    :param rate: the EMA rate (closer to 1 means slower).
+    """
+    for targ, src in zip(target_params, source_params):
+        targ.detach().mul_(rate).add_(src, alpha=1 - rate)
+
+def logger_file(path):
+    logger = logging.getLogger()
+    logger.setLevel(logging.DEBUG)
+    handler = logging.FileHandler(path,"w", encoding=None, delay="true")
+    handler.setLevel(logging.INFO)
+    formatter = logging.Formatter("%(message)s")
+    handler.setFormatter(formatter)
+    logger.addHandler(handler)
+    return logger
+
+
+def get_args_parser():
+    parser = argparse.ArgumentParser('MAR training with Diffusion Loss', add_help=False)
+    parser.add_argument('--batch_size', default=2, type=int,
+                        help='Batch size per GPU (effective batch size is batch_size * # gpus')
+    parser.add_argument('--epochs', default=2000, type=int)
+
+    # Model parameters
+    parser.add_argument('--model', default='mar_base', type=str, metavar='MODEL',
+                        help='Name of model to train')
+    parser.add_argument('--ckpt_path', default="pretrained_models/mar/city768.16.pth", type=str,
+                        help='model checkpoint path')
+
+    # VAE parameters
+    parser.add_argument('--img_size', default=768, type=int,
+                        help='images input size')
+    parser.add_argument('--vae_path', default="pretrained_models/vae/modelf16.ckpt", type=str,
+                        help='images input size')
+    parser.add_argument('--vae_embed_dim', default=16, type=int,
+                        help='vae output embedding dimension')
+    parser.add_argument('--vae_stride', default=16, type=int,
+                        help='tokenizer stride, default use KL16')
+    parser.add_argument('--patch_size', default=1, type=int,
+                        help='number of tokens to group as a patch.')
+    parser.add_argument('--config', default="ldm/config.yaml", type=str,
+                        help='vae model configuration file')
+
+    # Generation parameters
+    parser.add_argument('--num_iter', default=64, type=int,
+                        help='number of autoregressive iterations to generate an image')
+    parser.add_argument('--num_images', default=3000, type=int,
+                        help='number of images to generate')
+    parser.add_argument('--cfg', default=1.0, type=float, help="classifier-free guidance")
+    parser.add_argument('--cfg_schedule', default="linear", type=str)
+    parser.add_argument('--label_drop_prob', default=0.1, type=float)
+    parser.add_argument('--eval_freq', type=int, default=40, help='evaluation frequency')
+    parser.add_argument('--save_last_freq', type=int, default=5, help='save last frequency')
+    parser.add_argument('--online_eval', action='store_true')
+    parser.add_argument('--evaluate', action='store_true')
+    parser.add_argument('--eval_bsz', type=int, default=64, help='generation batch size')
+
+    # Optimizer parameters
+    parser.add_argument('--weight_decay', type=float, default=0.02,
+                        help='weight decay (default: 0.02)')
+
+    parser.add_argument('--grad_checkpointing', action='store_true')
+    parser.add_argument('--lr', type=float, default=None, metavar='LR',
+                        help='learning rate (absolute lr)')
+    parser.add_argument('--blr', type=float, default=1e-4, metavar='LR',
+                        help='base learning rate: absolute_lr = base_lr * total_batch_size / 256')
+    parser.add_argument('--min_lr', type=float, default=0., metavar='LR',
+                        help='lower lr bound for cyclic schedulers that hit 0')
+    parser.add_argument('--lr_schedule', type=str, default='constant',
+                        help='learning rate schedule')
+    parser.add_argument('--warmup_epochs', type=int, default=100, metavar='N',
+                        help='epochs to warmup LR')
+    parser.add_argument('--ema_rate', default=0.9999, type=float)
+
+    # MAR params
+    parser.add_argument('--mask_ratio_min', type=float, default=0.7,
+                        help='Minimum mask ratio')
+    parser.add_argument('--grad_clip', type=float, default=3.0,
+                        help='Gradient clip')
+    parser.add_argument('--attn_dropout', type=float, default=0.1,
+                        help='attention dropout')
+    parser.add_argument('--proj_dropout', type=float, default=0.1,
+                        help='projection dropout')
+    parser.add_argument('--buffer_size', type=int, default=64)
+
+    # Diffusion Loss params
+    parser.add_argument('--diffloss_d', type=int, default=6)
+    parser.add_argument('--diffloss_w', type=int, default=1024)
+    parser.add_argument('--num_sampling_steps', type=str, default="100")
+    parser.add_argument('--diffusion_batch_mul', type=int, default=4)
+    parser.add_argument('--temperature', default=1.0, type=float, help='diffusion loss sampling temperature')
+
+    # Dataset parameters
+    parser.add_argument('--output_dir', default='./output_dir',
+                        help='path where to save, empty for no saving')
+    parser.add_argument('--log_dir', default='./output_dir',
+                        help='path where to tensorboard log')
+    parser.add_argument('--device', default='cuda',
+                        help='device to use for training / testing')
+    parser.add_argument('--seed', default=1, type=int)
+    parser.add_argument('--resume', default=None,
+                        help='resume from checkpoint')
+
+    parser.add_argument('--start_epoch', default=0, type=int, metavar='N',
+                        help='start epoch')
+    parser.add_argument('--num_workers', default=10, type=int)
+    parser.add_argument('--pin_mem', action='store_true',
+                        help='Pin CPU memory in DataLoader for more efficient (sometimes) transfer to GPU.')
+    parser.add_argument('--no_pin_mem', action='store_false', dest='pin_mem')
+    parser.set_defaults(pin_mem=True)
+
+    # distributed training parameters
+    parser.add_argument('--world_size', default=1, type=int,
+                        help='number of distributed processes')
+    parser.add_argument('--local_rank', default=-1, type=int)
+    parser.add_argument('--dist_on_itp', action='store_true')
+    parser.add_argument('--dist_url', default='env://',
+                        help='url used to set up distributed training')
+
+    # caching latents
+    parser.add_argument('--use_cached', action='store_true', dest='use_cached',
+                        help='Use cached latents')
+    parser.set_defaults(use_cached=False)
+    parser.add_argument('--cached_path', default='', help='path to cached latents')
+
+    return parser
+
+
+def main(args):
+    misc.init_distributed_mode(args)
+
+    print('job dir: {}'.format(os.path.dirname(os.path.realpath(__file__))))
+    print("{}".format(args).replace(', ', ',\n'))
+
+    device = torch.device(args.device)
+
+    # fix the seed for reproducibility
+    seed = args.seed + misc.get_rank()
+    torch.manual_seed(seed)
+    np.random.seed(seed)
+
+    cudnn.benchmark = True
+
+    num_tasks = misc.get_world_size()
+    global_rank = misc.get_rank()
+
+    log_writer = None
+
+    # augmentation following DiT and ADM
+    transform_train = transforms.Compose([
+        transforms.ToTensor(),
+        transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])
+    ])
+
+    dataset_train = cityscapes.CityScapes('dataset/CityScapes/trainlist.txt', transform=transform_train, img_size=args.img_size)
+
+    sampler_train = torch.utils.data.DistributedSampler(
+        dataset_train, num_replicas=num_tasks, rank=global_rank, shuffle=True
+    )
+    print("Sampler_train = %s" % str(sampler_train))
+
+    data_loader_train = torch.utils.data.DataLoader(
+        dataset_train, sampler=sampler_train,
+        batch_size=args.batch_size,
+        num_workers=args.num_workers,
+        pin_memory=args.pin_mem,
+        drop_last=True,
+    )
+
+    # define the vae and mar model
+    with open(args.config, "r") as f:
+        config = yaml.safe_load(f)
+    args.ddconfig = config["ddconfig"]
+    print('cofig: ', config)
+
+    vae = AutoencoderKL(
+        ddconfig=args.ddconfig,
+        embed_dim=args.vae_embed_dim,
+        ckpt_path=args.vae_path
+    ).cuda().eval()
+
+    for param in vae.parameters():
+        param.requires_grad = False
+
+    model = mar.__dict__[args.model](
+        img_size=args.img_size,
+        vae_stride=args.vae_stride,
+        patch_size=args.patch_size,
+        vae_embed_dim=args.vae_embed_dim,
+        mask_ratio_min=args.mask_ratio_min,
+        label_drop_prob=args.label_drop_prob,
+        attn_dropout=args.attn_dropout,
+        proj_dropout=args.proj_dropout,
+        buffer_size=args.buffer_size,
+        diffloss_d=args.diffloss_d,
+        diffloss_w=args.diffloss_w,
+        num_sampling_steps=args.num_sampling_steps,
+        diffusion_batch_mul=args.diffusion_batch_mul,
+        grad_checkpointing=args.grad_checkpointing,
+    )
+
+    if args.ckpt_path:
+        checkpoint = torch.load(args.ckpt_path, map_location='cpu')
+        model.load_state_dict(checkpoint['model'])
+
+    print("Model = %s" % str(model))
+    # following timm: set wd as 0 for bias and norm layers
+    n_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
+    print("Number of trainable parameters: {}M".format(n_params / 1e6))
+
+    model.to(device)
+    model_without_ddp = model
+
+    eff_batch_size = args.batch_size * misc.get_world_size()
+
+    if args.lr is None:  # only base_lr is specified
+        args.lr = args.blr
+
+    print("base lr: %.2e" % args.blr)
+    print("actual lr: %.2e" % args.lr)
+    print("effective batch size: %d" % eff_batch_size)
+
+    if args.distributed:
+        model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.gpu])
+        model_without_ddp = model.module
+
+    # no weight decay on bias, norm layers, and diffloss MLP
+    param_groups = misc.add_weight_decay(model_without_ddp, args.weight_decay)
+    optimizer = torch.optim.AdamW(param_groups, lr=args.lr, betas=(0.9, 0.95))
+    print(optimizer)
+    loss_scaler = NativeScaler()
+
+    # resume training
+    if args.resume and glob.glob(os.path.join(args.output_dir, args.resume, 'checkpoint*.pth')):
+        try:
+            checkpoint = torch.load(sorted(glob.glob(os.path.join(args.output_dir, args.resume, 'checkpoint*.pth')))[-1], map_location='cpu')
+            model.load_state_dict(checkpoint['model'])
+        except:
+            checkpoint = torch.load(sorted(glob.glob(os.path.join(args.output_dir, args.resume, 'checkpoint*.pth')))[-2], map_location='cpu')
+            model.load_state_dict(checkpoint['model'])
+        state_dict = {key.replace("module.", ""): value for key, value in checkpoint['model'].items()}
+        model_without_ddp.load_state_dict(state_dict)
+        model_params = list(model_without_ddp.parameters())
+        ema_params = copy.deepcopy(model_params)
+        ema_state_dict = {key.replace("module.", ""): value for key, value in checkpoint['model_ema'].items()}
+        ema_params = [ema_state_dict[name].cuda() for name, _ in model_without_ddp.named_parameters()]
+        print("Resume checkpoint %s" % args.resume)
+        if 'optimizer' in checkpoint and 'epoch' in checkpoint:
+            optimizer.load_state_dict(checkpoint['optimizer'])
+            args.start_epoch = checkpoint['epoch'] + 1
+            if 'scaler' in checkpoint:
+                loss_scaler.load_state_dict(checkpoint['scaler'])
+            print("With optim & sched!")
+        del checkpoint
+
+        args.output_dir = os.path.join(args.output_dir, args.resume)
+
+        logger = logger_file(args.log_dir+'/'+args.resume+'.log')
+        if os.path.exists(args.log_dir+'/'+args.resume+'.log'):
+            with open(args.log_dir+'/'+args.resume+'.log', 'r') as infile:
+                for line in infile:
+                    logger.info(line.rstrip())
+        else:
+            logger.info("All the arguments")
+            for k, v in vars(args).items():
+                logger.info(f"{k}: {v}")
+            logger.info("\n\n Loss information")
+
+
+
+    else:
+        model_params = list(model_without_ddp.parameters())
+        ema_params = copy.deepcopy(model_params)
+        print("Training from scratch")
+        args.resume = datetime.datetime.now().strftime("%Y.%m.%d.%H.%M")        
+        args.output_dir = os.path.join(args.output_dir, args.resume)
+        Path(args.output_dir).mkdir(parents=True, exist_ok=True)
+
+        logger = logger_file(args.log_dir+'/'+args.resume+'.log')
+        logger.info("All the arguments")
+        for k, v in vars(args).items():
+            logger.info(f"{k}: {v}")
+        logger.info("\n\n Loss information")
+    
+
+    print(f"Start training for {args.epochs} epochs")
+    start_time = time.time()
+    for epoch in range(args.start_epoch, args.epochs):
+        if args.distributed:
+            data_loader_train.sampler.set_epoch(epoch)
+
+
+
+    for epoch in tqdm(range(args.start_epoch, args.epochs), desc="Training Progress"):
+        model.train(True)
+        metric_logger = misc.MetricLogger(delimiter="  ")
+        metric_logger.add_meter('lr', misc.SmoothedValue(window_size=1, fmt='{value:.6f}'))
+        header = 'Epoch: [{}]'.format(epoch)
+        print_freq = 20
+
+        optimizer.zero_grad()
+
+        for data_iter_step, (samples, labels, _) in enumerate(data_loader_train):
+            lr_sched.adjust_learning_rate(optimizer, data_iter_step / len(data_loader_train) + epoch, args)
+            samples = samples.to(device, non_blocking=True)
+            labels = labels.to(device, non_blocking=True)
+
+            with torch.no_grad():
+                posterior_x = vae.encode(samples)
+                posterior_y = vae.encode(labels)
+                x = posterior_x.sample().mul_(0.2325)
+                y = posterior_y.sample().mul_(0.2325)
+            with torch.cuda.amp.autocast():
+                loss = model(x,y)
+            loss_value = loss.item()
+            loss_scaler(loss, optimizer, clip_grad=args.grad_clip, parameters=model.parameters(), update_grad=True)
+            optimizer.zero_grad()
+            torch.cuda.synchronize()
+
+            update_ema(ema_params, model_params, rate=args.ema_rate)
+            metric_logger.update(loss=loss_value)
+
+            lr = optimizer.param_groups[0]["lr"]
+            metric_logger.update(lr=lr)
+
+            loss_value_reduce = misc.all_reduce_mean(loss_value)
+        metric_logger.synchronize_between_processes()
+        logger.info(f"epoch: {epoch:4d}, Averaged stats: {metric_logger}")
+        if (epoch+1)% args.save_last_freq == 0:
+            misc.save_model(args=args, model=model, model_without_ddp=model, optimizer=optimizer,
+                            loss_scaler=loss_scaler, epoch=epoch, ema_params=ema_params, epoch_name=str(epoch).zfill(5))
+
+    total_time = time.time() - start_time
+    total_time_str = str(datetime.timedelta(seconds=int(total_time)))
+    print('Training time {}'.format(total_time_str))
+
+
+if __name__ == '__main__':
+    args = get_args_parser()
+    args = args.parse_args()
+    Path(args.output_dir).mkdir(parents=True, exist_ok=True)
+    Path(args.log_dir).mkdir(parents=True, exist_ok=True)
+    main(args)

validation.ipynb

@@ -0,0 +1,364 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "c524f796-e657-4a59-abcf-540531a38995",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "%run get_parser.py"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "4c1cf01e-8229-4d28-bcb2-01c07fa641c2",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "import requests\n",
+    "\n",
+    "# Define URLs and file paths\n",
+    "files_to_download = {\n",
+    "    \"https://huggingface.co/mahmed10/CAM-Seg/resolve/main/pretrained_models/vae/modelf16.ckpt\":\n",
+    "        \"pretrained_models/vae/modelf16.ckpt\",\n",
+    "    \"https://huggingface.co/mahmed10/CAM-Seg/resolve/main/pretrained_models/mar/city768.16.pth\":\n",
+    "        \"pretrained_models/mar/city768.16.pth\"\n",
+    "}\n",
+    "\n",
+    "for url, path in files_to_download.items():\n",
+    "    os.makedirs(os.path.dirname(path), exist_ok=True)\n",
+    "    \n",
+    "    if os.path.exists(path):\n",
+    "        print(f\"File already exists: {path} — skipping download.\")\n",
+    "        continue\n",
+    "\n",
+    "    print(f\"Downloading from {url}...\")\n",
+    "    response = requests.get(url, stream=True)\n",
+    "    if response.status_code == 200:\n",
+    "        with open(path, 'wb') as f:\n",
+    "            for chunk in response.iter_content(chunk_size=8192):\n",
+    "                f.write(chunk)\n",
+    "        print(f\"Saved to {path}\")\n",
+    "    else:\n",
+    "        print(f\"Failed to download from {url}, status code {response.status_code}\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "id": "3a7ac93b-1cbc-45f3-8ec5-8e8257a39786",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "import numpy as np\n",
+    "from tqdm import tqdm\n",
+    "from PIL import Image\n",
+    "import yaml\n",
+    "import math\n",
+    "\n",
+    "import torch\n",
+    "import torch.backends.cudnn as cudnn\n",
+    "import torchvision.transforms as transforms\n",
+    "\n",
+    "from data import cityscapes\n",
+    "import util.misc as misc\n",
+    "\n",
+    "from models.vae import AutoencoderKL\n",
+    "from models import mar"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "id": "e2bde6fd-9b39-40fd-8d4d-d0a5f9c8217a",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "def mask_by_order(mask_len, order, bsz, seq_len):\n",
+    "    masking = torch.zeros(bsz, seq_len).cuda()\n",
+    "    masking = torch.scatter(masking, dim=-1, index=order[:, :mask_len.long()], src=torch.ones(bsz, seq_len).cuda()).bool()\n",
+    "    return masking\n",
+    "\n",
+    "def fast_hist(pred, label, n):\n",
+    "    k = (label >= 0) & (label < n)\n",
+    "    bin_count = np.bincount(\n",
+    "        n * label[k].astype(int) + pred[k], minlength=n ** 2)\n",
+    "    return bin_count[:n ** 2].reshape(n, n)\n",
+    "\n",
+    "color_pallete = np.round(np.array([\n",
+    "            0, 0, 0,\n",
+    "            128, 64, 128,\n",
+    "            244, 35, 232,\n",
+    "            70, 70, 70,\n",
+    "            102, 102, 156,\n",
+    "            190, 153, 153,\n",
+    "            153, 153, 153,\n",
+    "            250, 170, 30,\n",
+    "            220, 220, 0,\n",
+    "            107, 142, 35,\n",
+    "            152, 251, 152,\n",
+    "            0, 130, 180,\n",
+    "            220, 20, 60,\n",
+    "            255, 0, 0,\n",
+    "            0, 0, 142,\n",
+    "            0, 0, 70,\n",
+    "            0, 60, 100,\n",
+    "            0, 80, 100,\n",
+    "            0, 0, 230,\n",
+    "            119, 11, 32,\n",
+    "        ])/255.0, 4)\n",
+    "\n",
+    "color_pallete = color_pallete.reshape(-1, 3)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "id": "c189ac7b-ccff-4745-af56-460ec88770b4",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "device = torch.device(args.device)\n",
+    "device = torch.device('cuda:0')\n",
+    "args.batch_size = 1\n",
+    "\n",
+    "# fix the seed for reproducibility\n",
+    "seed = args.seed + misc.get_rank()\n",
+    "torch.manual_seed(seed)\n",
+    "np.random.seed(seed)\n",
+    "\n",
+    "cudnn.benchmark = True\n",
+    "\n",
+    "num_tasks = misc.get_world_size()\n",
+    "global_rank = misc.get_rank()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "id": "28d13453-a3ac-4d2e-8906-0c179e85c2f9",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "transform_train = transforms.Compose([\n",
+    "    transforms.ToTensor(),\n",
+    "    transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])\n",
+    "])\n",
+    "\n",
+    "dataset_train = cityscapes.CityScapes('dataset/CityScapes/vallist.txt', data_set= 'val', transform=transform_train, seed=args.seed, img_size=args.img_size)\n",
+    "# dataset_train = umbc.UMBC('dataset/UMBC/all.txt', data_set= 'val', transform=transform_train, seed=args.seed, img_size=args.img_size)\n",
+    "# dataset_train = acdc.ACDC('dataset/ACDC/vallist_fog.txt', data_set= 'val', transform=transform_train, seed=args.seed, img_size=args.img_size)\n",
+    "# dataset_train = semantickitti.SemanticKITTI('dataset/SemanticKitti/vallist.txt', data_set= 'val', transform=transform_train, seed=args.seed, img_size=args.img_size)\n",
+    "\n",
+    "\n",
+    "sampler_train = torch.utils.data.DistributedSampler(dataset_train, num_replicas=1, rank=0, shuffle=False)\n",
+    "\n",
+    "data_loader_train = torch.utils.data.DataLoader(\n",
+    "    dataset_train, sampler=sampler_train,\n",
+    "    batch_size=args.batch_size,\n",
+    "    num_workers=args.num_workers,\n",
+    "    pin_memory=args.pin_mem,\n",
+    "    drop_last=True,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "2e22d231-02db-4586-b489-01a97314aed9",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "vae = AutoencoderKL(\n",
+    "    ddconfig=args.ddconfig,\n",
+    "    embed_dim=args.vae_embed_dim,\n",
+    "    ckpt_path=args.vae_path\n",
+    ").to(device).eval()\n",
+    "\n",
+    "for param in vae.parameters():\n",
+    "    param.requires_grad = False\n",
+    "    \n",
+    "model = mar.mar_base(\n",
+    "    img_size=args.img_size,\n",
+    "    vae_stride=args.vae_stride,\n",
+    "    patch_size=args.patch_size,\n",
+    "    vae_embed_dim=args.vae_embed_dim,\n",
+    "    mask_ratio_min=args.mask_ratio_min,\n",
+    "    label_drop_prob=args.label_drop_prob,\n",
+    "    attn_dropout=args.attn_dropout,\n",
+    "    proj_dropout=args.proj_dropout,\n",
+    "    buffer_size=args.buffer_size,\n",
+    "    diffloss_d=args.diffloss_d,\n",
+    "    diffloss_w=args.diffloss_w,\n",
+    "    num_sampling_steps=args.num_sampling_steps,\n",
+    "    diffusion_batch_mul=args.diffusion_batch_mul,\n",
+    "    grad_checkpointing=args.grad_checkpointing,\n",
+    ")\n",
+    "\n",
+    "n_params = sum(p.numel() for p in model.parameters() if p.requires_grad)\n",
+    "print(\"Number of trainable parameters: {}M\".format(n_params / 1e6))\n",
+    "\n",
+    "\n",
+    "checkpoint = torch.load(args.ckpt_path, map_location='cpu')\n",
+    "model.load_state_dict(checkpoint['model'])\n",
+    "model.to(device)\n",
+    "\n",
+    "eff_batch_size = args.batch_size * misc.get_world_size()\n",
+    "\n",
+    "print(\"effective batch size: %d\" % eff_batch_size)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "id": "4c83c0eb-35a5-4241-b869-d52eb6cd31e0",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Training Progress: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 500/500 [13:11<00:00,  1.58s/it]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "road        :  98.06\n",
+      "sidewalk    :  86.32\n",
+      "building    :  89.23\n",
+      "wall        :  47.44\n",
+      "fence       :  43.78\n",
+      "pole        :  60.14\n",
+      "tlight      :  63.16\n",
+      "tsign       :  82.48\n",
+      "vtation     :  92.72\n",
+      "terrain     :  80.45\n",
+      "sky         :  95.99\n",
+      "person      :  70.83\n",
+      "rider       :  64.25\n",
+      "car         :  94.06\n",
+      "truck       :  44.90\n",
+      "bus         :  66.81\n",
+      "train       :  44.04\n",
+      "motorcycle  :  47.34\n",
+      "bicycle     :  62.50\n",
+      "Avg Pre     :  70.24\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\n"
+     ]
+    }
+   ],
+   "source": [
+    "hist = []\n",
+    "model.eval()\n",
+    "for data_iter_step, (samples, labels, path) in enumerate(tqdm(data_loader_train, desc=\"Training Progress\")):\n",
+    "    samples = samples.to(device, non_blocking=True)\n",
+    "    labels = labels.to(device, non_blocking=True)\n",
+    "\n",
+    "    with torch.no_grad():\n",
+    "        posterior_x = vae.encode(samples)\n",
+    "        posterior_y = vae.encode(labels)\n",
+    "        x = posterior_x.sample().mul_(0.2325)\n",
+    "        y = posterior_y.sample().mul_(0.2325)\n",
+    "        x = model.patchify(x)\n",
+    "        y = model.patchify(y)\n",
+    "        gt_latents = y.clone().detach()\n",
+    "    cfg_iter = 1.0\n",
+    "    temperature = 1.0\n",
+    "    mask_actual = torch.cat([torch.zeros(args.batch_size, model.seq_len), torch.ones(args.batch_size, model.seq_len)], dim=1).cuda()\n",
+    "    tokens = torch.zeros(args.batch_size, model.seq_len, model.token_embed_dim).cuda()\n",
+    "\n",
+    "    with torch.no_grad():\n",
+    "        x1 = model.forward_mae_encoder(x, mask_actual, tokens)\n",
+    "        z = model.forward_mae_decoder(x1, mask_actual)\n",
+    "    z = z[0]\n",
+    "    sampled_token_latent = model.diffloss.sample(z, temperature, cfg_iter)\n",
+    "\n",
+    "    tokens[0] = sampled_token_latent[model.seq_len:]\n",
+    "    tokens = model.unpatchify(tokens)\n",
+    "    \n",
+    "    sampled_images = vae.decode(tokens / 0.2325)\n",
+    "    \n",
+    "    image_tensor = labels[0] \n",
+    "    image_tensor = image_tensor * 0.5 + 0.5\n",
+    "    gt_np = image_tensor.permute(1, 2, 0).cpu().numpy()\n",
+    "    H, W, _ = gt_np.shape\n",
+    "    pixels = gt_np.reshape(-1, 3)\n",
+    "    distances = np.linalg.norm(pixels[:, None, :] - color_pallete[None, :, :], axis=2)\n",
+    "    output = np.argmin(distances, axis=1)\n",
+    "    gt = output.reshape(H, W)\n",
+    "    \n",
+    "    image_tensor = sampled_images[0]\n",
+    "    image_tensor = image_tensor * 0.5 + 0.5 \n",
+    "    ss_np = image_tensor.permute(1, 2, 0).cpu().numpy()\n",
+    "    H, W, _ = ss_np.shape\n",
+    "    pixels = ss_np.reshape(-1, 3)\n",
+    "    distances = np.linalg.norm(pixels[:, None, :] - color_pallete[None, :, :], axis=2)\n",
+    "    output = np.argmin(distances, axis=1)\n",
+    "    output = output.reshape(H, W)\n",
+    "    \n",
+    "    hist.append(fast_hist(output.reshape(-1), gt.reshape(-1), 20))\n",
+    "\n",
+    "cm = np.sum(hist, axis=0)\n",
+    "\n",
+    "epsilon = 1e-10\n",
+    "class_precision = np.diag(cm[1:,1:]) / (np.sum(cm[1:,1:], axis=0) + epsilon)\n",
+    "class_names = ['road', 'sidewalk', 'building', 'wall', 'fence', 'pole', 'tlight', 'tsign', \n",
+    "               'vtation', 'terrain', 'sky', 'person', 'rider', 'car', 'truck', 'bus', 'train', \n",
+    "               'motorcycle', 'bicycle']\n",
+    "\n",
+    "for i in range(len(class_names)):\n",
+    "    print(f\"{class_names[i]:<12}: {class_precision[i]*100:6.2f}\")\n",
+    "average_precision = np.mean(class_precision)\n",
+    "print(f\"{'Avg Pre':<12}: {average_precision*100:6.2f}\")"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.8.10"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}