diff --git a/README.md b/README.md
index 8fb44f7c6c6045c9106e20accfa1d3e6df0dc4f0..4e741fb61c446b6e724d747412c633b173e1a1c6 100644
--- a/README.md
+++ b/README.md
@@ -1,8 +1,8 @@
---
title: LiDAR Diffusion
-emoji: 📚
-colorFrom: blue
-colorTo: green
+emoji: 🚙🛞🚨
+colorFrom: green
+colorTo: indigo
sdk: gradio
sdk_version: 4.26.0
app_file: app.py
diff --git a/app.py b/app.py
index d860a45d46dcefd21abcdb62e4f2b5237a938485..3eb1fe7039b19818959a83ab682e32b92e0d730c 100644
--- a/app.py
+++ b/app.py
@@ -1,9 +1,81 @@
import gradio as gr
+import spaces
+import tempfile
+import os
+import torch
+import numpy as np
+from matplotlib.colors import LinearSegmentedColormap
+from app_config import CSS, TITLE, DESCRIPTION, DEVICE
+import sample_cond
-def greet(name):
- return "Hello " + name + "!!"
+model = sample_cond.load_model()
-iface = gr.Interface(fn=greet, inputs="text", outputs="text")
-iface.launch()
+def create_custom_colormap():
+ colors = [(0, 1, 0), (0, 1, 1), (0, 0, 1), (1, 0, 1), (1, 1, 0)]
+ positions = [0, 0.38, 0.6, 0.7, 1]
+
+ custom_cmap = LinearSegmentedColormap.from_list('custom_colormap', list(zip(positions, colors)), N=256)
+ return custom_cmap
+
+
+def colorize_depth(depth, log_scale):
+ if log_scale:
+ depth = ((np.log2((depth / 255.) * 56. + 1) / 5.84) * 255.).astype(np.uint8)
+ mask = depth == 0
+ colormap = create_custom_colormap()
+ rgb = colormap(depth)[:, :, :3]
+ rgb[mask] = 0.
+ return rgb
+
+
+@spaces.GPU
+@torch.no_grad()
+def generate_lidar(model, cond):
+ img, pcd = sample_cond.sample(model, cond)
+ return img, pcd
+
+
+def load_camera(image):
+ split_per_view = 4
+ camera = np.array(image).astype(np.float32) / 255.
+ camera = camera.transpose(2, 0, 1)
+ camera_list = np.split(camera, split_per_view, axis=2) # split into n chunks as different views
+ camera_cond = torch.from_numpy(np.stack(camera_list, axis=0)).unsqueeze(0).to(DEVICE)
+ return camera_cond
+
+
+with gr.Blocks(css=CSS) as demo:
+ gr.Markdown(TITLE)
+ gr.Markdown(DESCRIPTION)
+ gr.Markdown("### Camera-to-LiDAR Demo")
+ # gr.Markdown("You can slide the output to compare the depth prediction with input image")
+
+ with gr.Row():
+ input_image = gr.Image(label="Input Image", type='numpy', elem_id='img-display-input')
+ output_image = gr.Image(label="Range Map", elem_id='img-display-output')
+ raw_file = gr.File(label="Point Cloud (.txt file). Can be viewed through Meshlab")
+ submit = gr.Button("Submit")
+
+ def on_submit(image):
+ cond = load_camera(image)
+ img, pcd = generate_lidar(model, cond)
+
+ tmp = tempfile.NamedTemporaryFile(suffix='.txt', delete=False)
+ pcd.save(tmp.name)
+
+ rgb_img = colorize_depth(img, log_scale=True)
+
+ return [rgb_img, tmp.name]
+
+ submit.click(on_submit, inputs=[input_image], outputs=[output_image, raw_file])
+
+ example_files = sorted(os.listdir('cam_examples'))
+ example_files = [os.path.join('cam_examples', filename) for filename in example_files]
+ examples = gr.Examples(examples=example_files, inputs=[input_image], outputs=[output_image, raw_file],
+ fn=on_submit, cache_examples=True)
+
+
+if __name__ == '__main__':
+ demo.queue().launch()
diff --git a/app_config.py b/app_config.py
new file mode 100644
index 0000000000000000000000000000000000000000..505160872e64d4fbdab7e564d8388a0afe647eed
--- /dev/null
+++ b/app_config.py
@@ -0,0 +1,17 @@
+import torch
+
+CSS = """
+#img-display-container {
+ max-height: 100vh;
+ }
+#img-display-input {
+ max-height: 80vh;
+ }
+#img-display-output {
+ max-height: 80vh;
+ }
+"""
+DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
+TITLE = "# LiDAR Diffusion"
+DESCRIPTION = """Official demo for **LiDAR Diffusion: Towards Realistic Scene Generation with LiDAR Diffusion Models**.
+Please refer to our [paper](https://arxiv.org/abs/2404.00815), [project page](https://lidar-diffusion.github.io/), or [github](https://github.com/hancyran/LiDAR-Diffusion) for more details."""
diff --git a/cam_examples/conditioning_000011.png b/cam_examples/conditioning_000011.png
new file mode 100644
index 0000000000000000000000000000000000000000..f8e13c1e65b515ee85ababb900fcb5316ac37086
Binary files /dev/null and b/cam_examples/conditioning_000011.png differ
diff --git a/cam_examples/conditioning_000153.png b/cam_examples/conditioning_000153.png
new file mode 100644
index 0000000000000000000000000000000000000000..81f0c60659b3f31b2c01cfc750780c72cd344f9f
Binary files /dev/null and b/cam_examples/conditioning_000153.png differ
diff --git a/cam_examples/conditioning_000354.png b/cam_examples/conditioning_000354.png
new file mode 100644
index 0000000000000000000000000000000000000000..02b534675afac8bb83cf944c9098884fb6a74298
Binary files /dev/null and b/cam_examples/conditioning_000354.png differ
diff --git a/cam_examples/conditioning_000555.png b/cam_examples/conditioning_000555.png
new file mode 100644
index 0000000000000000000000000000000000000000..0d01f5c9be3b6118fad8fba41e4ffd11837d5676
Binary files /dev/null and b/cam_examples/conditioning_000555.png differ
diff --git a/cam_examples/conditioning_001026.png b/cam_examples/conditioning_001026.png
new file mode 100644
index 0000000000000000000000000000000000000000..266b1de4f9cd659d15ac6d176b3411753cd42159
Binary files /dev/null and b/cam_examples/conditioning_001026.png differ
diff --git a/data/config/semantic-kitti.yaml b/data/config/semantic-kitti.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..628106553b4b964920c825af42a53b4e39b73cfd
--- /dev/null
+++ b/data/config/semantic-kitti.yaml
@@ -0,0 +1,211 @@
+# This file is covered by the LICENSE file in the root of this project.
+labels:
+ 0 : "unlabeled"
+ 1 : "outlier"
+ 10: "car"
+ 11: "bicycle"
+ 13: "bus"
+ 15: "motorcycle"
+ 16: "on-rails"
+ 18: "truck"
+ 20: "other-vehicle"
+ 30: "person"
+ 31: "bicyclist"
+ 32: "motorcyclist"
+ 40: "road"
+ 44: "parking"
+ 48: "sidewalk"
+ 49: "other-ground"
+ 50: "building"
+ 51: "fence"
+ 52: "other-structure"
+ 60: "lane-marking"
+ 70: "vegetation"
+ 71: "trunk"
+ 72: "terrain"
+ 80: "pole"
+ 81: "traffic-sign"
+ 99: "other-object"
+ 252: "moving-car"
+ 253: "moving-bicyclist"
+ 254: "moving-person"
+ 255: "moving-motorcyclist"
+ 256: "moving-on-rails"
+ 257: "moving-bus"
+ 258: "moving-truck"
+ 259: "moving-other-vehicle"
+color_map: # bgr
+ 0 : [0, 0, 0]
+ 1 : [0, 0, 255]
+ 10: [245, 150, 100]
+ 11: [245, 230, 100]
+ 13: [250, 80, 100]
+ 15: [150, 60, 30]
+ 16: [255, 0, 0]
+ 18: [180, 30, 80]
+ 20: [255, 0, 0]
+ 30: [30, 30, 255]
+ 31: [200, 40, 255]
+ 32: [90, 30, 150]
+ 40: [255, 0, 255]
+ 44: [255, 150, 255]
+ 48: [75, 0, 75]
+ 49: [75, 0, 175]
+ 50: [0, 200, 255]
+ 51: [50, 120, 255]
+ 52: [0, 150, 255]
+ 60: [170, 255, 150]
+ 70: [0, 175, 0]
+ 71: [0, 60, 135]
+ 72: [80, 240, 150]
+ 80: [150, 240, 255]
+ 81: [0, 0, 255]
+ 99: [255, 255, 50]
+ 252: [245, 150, 100]
+ 256: [255, 0, 0]
+ 253: [200, 40, 255]
+ 254: [30, 30, 255]
+ 255: [90, 30, 150]
+ 257: [250, 80, 100]
+ 258: [180, 30, 80]
+ 259: [255, 0, 0]
+content: # as a ratio with the total number of points
+ 0: 0.018889854628292943
+ 1: 0.0002937197336781505
+ 10: 0.040818519255974316
+ 11: 0.00016609538710764618
+ 13: 2.7879693665067774e-05
+ 15: 0.00039838616015114444
+ 16: 0.0
+ 18: 0.0020633612104619787
+ 20: 0.0016218197275284021
+ 30: 0.00017698551338515307
+ 31: 1.1065903904919655e-08
+ 32: 5.532951952459828e-09
+ 40: 0.1987493871255525
+ 44: 0.014717169549888214
+ 48: 0.14392298360372
+ 49: 0.0039048553037472045
+ 50: 0.1326861944777486
+ 51: 0.0723592229456223
+ 52: 0.002395131480328884
+ 60: 4.7084144280367186e-05
+ 70: 0.26681502148037506
+ 71: 0.006035012012626033
+ 72: 0.07814222006271769
+ 80: 0.002855498193863172
+ 81: 0.0006155958086189918
+ 99: 0.009923127583046915
+ 252: 0.001789309418528068
+ 253: 0.00012709999297008662
+ 254: 0.00016059776092534436
+ 255: 3.745553104802113e-05
+ 256: 0.0
+ 257: 0.00011351574470342043
+ 258: 0.00010157861367183268
+ 259: 4.3840131989471124e-05
+# classes that are indistinguishable from single scan or inconsistent in
+# ground truth are mapped to their closest equivalent
+learning_map:
+ 0 : 0 # "unlabeled"
+ 1 : 0 # "outlier" mapped to "unlabeled" --------------------------mapped
+ 10: 1 # "car"
+ 11: 2 # "bicycle"
+ 13: 5 # "bus" mapped to "other-vehicle" --------------------------mapped
+ 15: 3 # "motorcycle"
+ 16: 5 # "on-rails" mapped to "other-vehicle" ---------------------mapped
+ 18: 4 # "truck"
+ 20: 5 # "other-vehicle"
+ 30: 6 # "person"
+ 31: 7 # "bicyclist"
+ 32: 8 # "motorcyclist"
+ 40: 9 # "road"
+ 44: 10 # "parking"
+ 48: 11 # "sidewalk"
+ 49: 12 # "other-ground"
+ 50: 13 # "building"
+ 51: 14 # "fence"
+ 52: 0 # "other-structure" mapped to "unlabeled" ------------------mapped
+ 60: 9 # "lane-marking" to "road" ---------------------------------mapped
+ 70: 15 # "vegetation"
+ 71: 16 # "trunk"
+ 72: 17 # "terrain"
+ 80: 18 # "pole"
+ 81: 19 # "traffic-sign"
+ 99: 0 # "other-object" to "unlabeled" ----------------------------mapped
+ 252: 1 # "moving-car" to "car" ------------------------------------mapped
+ 253: 7 # "moving-bicyclist" to "bicyclist" ------------------------mapped
+ 254: 6 # "moving-person" to "person" ------------------------------mapped
+ 255: 8 # "moving-motorcyclist" to "motorcyclist" ------------------mapped
+ 256: 5 # "moving-on-rails" mapped to "other-vehicle" --------------mapped
+ 257: 5 # "moving-bus" mapped to "other-vehicle" -------------------mapped
+ 258: 4 # "moving-truck" to "truck" --------------------------------mapped
+ 259: 5 # "moving-other"-vehicle to "other-vehicle" ----------------mapped
+learning_map_inv: # inverse of previous map
+ 0: 0 # "unlabeled", and others ignored
+ 1: 10 # "car"
+ 2: 11 # "bicycle"
+ 3: 15 # "motorcycle"
+ 4: 18 # "truck"
+ 5: 20 # "other-vehicle"
+ 6: 30 # "person"
+ 7: 31 # "bicyclist"
+ 8: 32 # "motorcyclist"
+ 9: 40 # "road"
+ 10: 44 # "parking"
+ 11: 48 # "sidewalk"
+ 12: 49 # "other-ground"
+ 13: 50 # "building"
+ 14: 51 # "fence"
+ 15: 70 # "vegetation"
+ 16: 71 # "trunk"
+ 17: 72 # "terrain"
+ 18: 80 # "pole"
+ 19: 81 # "traffic-sign"
+learning_ignore: # Ignore classes
+ 0: True # "unlabeled", and others ignored
+ 1: False # "car"
+ 2: False # "bicycle"
+ 3: False # "motorcycle"
+ 4: False # "truck"
+ 5: False # "other-vehicle"
+ 6: False # "person"
+ 7: False # "bicyclist"
+ 8: False # "motorcyclist"
+ 9: False # "road"
+ 10: False # "parking"
+ 11: False # "sidewalk"
+ 12: False # "other-ground"
+ 13: False # "building"
+ 14: False # "fence"
+ 15: False # "vegetation"
+ 16: False # "trunk"
+ 17: False # "terrain"
+ 18: False # "pole"
+ 19: False # "traffic-sign"
+split: # sequence numbers
+ train:
+ - 0
+ - 1
+ - 2
+ - 3
+ - 4
+ - 5
+ - 6
+ - 7
+ - 9
+ - 10
+ valid:
+ - 8
+ test:
+ - 11
+ - 12
+ - 13
+ - 14
+ - 15
+ - 16
+ - 17
+ - 18
+ - 19
+ - 20
+ - 21
diff --git a/lidm/data/__init__.py b/lidm/data/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/lidm/data/annotated_dataset.py b/lidm/data/annotated_dataset.py
new file mode 100644
index 0000000000000000000000000000000000000000..a81bf6fb186a145b1a654e23291329683d37d918
--- /dev/null
+++ b/lidm/data/annotated_dataset.py
@@ -0,0 +1,48 @@
+from pathlib import Path
+from typing import Optional, List, Dict, Union, Any
+import warnings
+
+from torch.utils.data import Dataset
+
+from .conditional_builder.objects_bbox import ObjectsBoundingBoxConditionalBuilder
+from .conditional_builder.objects_center_points import ObjectsCenterPointsConditionalBuilder
+
+
+class Annotated3DObjectsDataset(Dataset):
+ def __init__(self, min_objects_per_image: int,
+ max_objects_per_image: int, no_tokens: int, num_beams: int, cats: List[str],
+ cat_blacklist: Optional[List[str]] = None, **kwargs):
+ self.min_objects_per_image = min_objects_per_image
+ self.max_objects_per_image = max_objects_per_image
+ self.no_tokens = no_tokens
+ self.num_beams = num_beams
+
+ self.categories = [c for c in cats if c not in cat_blacklist] if cat_blacklist is not None else cats
+ self._conditional_builders = None
+
+ @property
+ def no_classes(self) -> int:
+ return len(self.categories)
+
+ @property
+ def conditional_builders(self) -> ObjectsCenterPointsConditionalBuilder:
+ # cannot set this up in init because no_classes is only known after loading data in init of superclass
+ if self._conditional_builders is None:
+ self._conditional_builders = {
+ 'center': ObjectsCenterPointsConditionalBuilder(
+ self.no_classes,
+ self.max_objects_per_image,
+ self.no_tokens,
+ self.num_beams
+ ),
+ 'bbox': ObjectsBoundingBoxConditionalBuilder(
+ self.no_classes,
+ self.max_objects_per_image,
+ self.no_tokens,
+ self.num_beams
+ )
+ }
+ return self._conditional_builders
+
+ def get_textual_label_for_category_id(self, category_id: int) -> str:
+ return self.categories[category_id]
diff --git a/lidm/data/base.py b/lidm/data/base.py
new file mode 100644
index 0000000000000000000000000000000000000000..ea9986df81f98aec686be929c3b620f660d02a92
--- /dev/null
+++ b/lidm/data/base.py
@@ -0,0 +1,121 @@
+import pdb
+from abc import abstractmethod
+from functools import partial
+
+import PIL
+import numpy as np
+from PIL import Image
+
+import torchvision.transforms.functional as TF
+from torch.utils.data import Dataset, IterableDataset
+
+from ..utils.aug_utils import get_lidar_transform, get_camera_transform, get_anno_transform
+
+
+class DatasetBase(Dataset):
+ def __init__(self, data_root, split, dataset_config, aug_config, return_pcd=False, condition_key=None,
+ scale_factors=None, degradation=None, **kwargs):
+ self.data_root = data_root
+ self.split = split
+ self.data = []
+ self.aug_config = aug_config
+
+ self.img_size = dataset_config.size
+ self.fov = dataset_config.fov
+ self.depth_range = dataset_config.depth_range
+ self.filtered_map_cats = dataset_config.filtered_map_cats
+ self.depth_scale = dataset_config.depth_scale
+ self.log_scale = dataset_config.log_scale
+
+ if self.log_scale:
+ self.depth_thresh = (np.log2(1./255. + 1) / self.depth_scale) * 2. - 1 + 1e-6
+ else:
+ self.depth_thresh = (1./255. / self.depth_scale) * 2. - 1 + 1e-6
+ self.return_pcd = return_pcd
+
+ if degradation is not None and scale_factors is not None:
+ scaled_img_size = (int(self.img_size[0] / scale_factors[0]), int(self.img_size[1] / scale_factors[1]))
+ degradation_fn = {
+ "pil_nearest": PIL.Image.NEAREST,
+ "pil_bilinear": PIL.Image.BILINEAR,
+ "pil_bicubic": PIL.Image.BICUBIC,
+ "pil_box": PIL.Image.BOX,
+ "pil_hamming": PIL.Image.HAMMING,
+ "pil_lanczos": PIL.Image.LANCZOS,
+ }[degradation]
+ self.degradation_transform = partial(TF.resize, size=scaled_img_size, interpolation=degradation_fn)
+ else:
+ self.degradation_transform = None
+ self.condition_key = condition_key
+
+ self.lidar_transform = get_lidar_transform(aug_config, split)
+ self.anno_transform = get_anno_transform(aug_config, split) if condition_key in ['bbox', 'center'] else None
+ self.view_transform = get_camera_transform(aug_config, split) if condition_key in ['camera'] else None
+
+ self.prepare_data()
+
+ def prepare_data(self):
+ raise NotImplementedError
+
+ def process_scan(self, range_img):
+ range_img = np.where(range_img < 0, 0, range_img)
+
+ if self.log_scale:
+ # log scale
+ range_img = np.log2(range_img + 0.0001 + 1)
+
+ range_img = range_img / self.depth_scale
+ range_img = range_img * 2. - 1.
+
+ range_img = np.clip(range_img, -1, 1)
+ range_img = np.expand_dims(range_img, axis=0)
+
+ # mask
+ range_mask = np.ones_like(range_img)
+ range_mask[range_img < self.depth_thresh] = -1
+
+ return range_img, range_mask
+
+ @staticmethod
+ def load_lidar_sweep(*args, **kwargs):
+ raise NotImplementedError
+
+ @staticmethod
+ def load_semantic_map(*args, **kwargs):
+ raise NotImplementedError
+
+ @staticmethod
+ def load_camera(*args, **kwargs):
+ raise NotImplementedError
+
+ @staticmethod
+ def load_annotation(*args, **kwargs):
+ raise NotImplementedError
+
+ def __len__(self):
+ return len(self.data)
+
+ def __getitem__(self, idx):
+ example = dict()
+ return example
+
+
+class Txt2ImgIterableBaseDataset(IterableDataset):
+ """
+ Define an interface to make the IterableDatasets for text2img data chainable
+ """
+ def __init__(self, num_records=0, valid_ids=None, size=256):
+ super().__init__()
+ self.num_records = num_records
+ self.valid_ids = valid_ids
+ self.sample_ids = valid_ids
+ self.size = size
+
+ print(f'{self.__class__.__name__} dataset contains {self.__len__()} examples.')
+
+ def __len__(self):
+ return self.num_records
+
+ @abstractmethod
+ def __iter__(self):
+ pass
\ No newline at end of file
diff --git a/lidm/data/conditional_builder/__init__.py b/lidm/data/conditional_builder/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/lidm/data/conditional_builder/objects_bbox.py b/lidm/data/conditional_builder/objects_bbox.py
new file mode 100644
index 0000000000000000000000000000000000000000..982da90403dbcd7bcada38651cf13efdfb89bc3d
--- /dev/null
+++ b/lidm/data/conditional_builder/objects_bbox.py
@@ -0,0 +1,53 @@
+from itertools import cycle
+from typing import List, Tuple, Callable, Optional
+
+from PIL import Image as pil_image, ImageDraw as pil_img_draw, ImageFont
+from more_itertools.recipes import grouper
+from torch import LongTensor, Tensor
+
+from ..helper_types import BoundingBox, Annotation
+from .objects_center_points import ObjectsCenterPointsConditionalBuilder, convert_pil_to_tensor
+from .utils import COLOR_PALETTE, WHITE, GRAY_75, BLACK, additional_parameters_string, \
+ pad_list, get_plot_font_size, absolute_bbox
+
+
+class ObjectsBoundingBoxConditionalBuilder(ObjectsCenterPointsConditionalBuilder):
+ @property
+ def object_descriptor_length(self) -> int:
+ return 3 # 3/5: object_representation (1) + corners (2/4)
+
+ def _make_object_descriptors(self, annotations: List[Annotation]) -> List[Tuple[int, ...]]:
+ object_tuples = [
+ (self.object_representation(ann), *self.token_pair_from_bbox(ann.bbox))
+ for ann in annotations
+ ]
+ object_tuples = pad_list(object_tuples, self.empty_tuple, self.no_max_objects)
+ return object_tuples
+
+ def inverse_build(self, conditional: LongTensor) -> Tuple[List[Tuple[int, BoundingBox]], Optional[BoundingBox]]:
+ conditional_list = conditional.tolist()
+ object_triples = grouper(conditional_list, 3)
+ assert conditional.shape[0] == self.embedding_dim
+ return [(object_triple[0], self.bbox_from_token_pair(object_triple[1], object_triple[2])) for object_triple in object_triples if object_triple[0] != self.none], None
+
+ def plot(self, conditional: LongTensor, label_for_category_no: Callable[[int], str], figure_size: Tuple[int, int],
+ line_width: int = 3, font_size: Optional[int] = None) -> Tensor:
+ plot = pil_image.new('RGB', figure_size, WHITE)
+ draw = pil_img_draw.Draw(plot)
+ # font = ImageFont.truetype(
+ # "/usr/share/fonts/truetype/lato/Lato-Regular.ttf",
+ # size=get_plot_font_size(font_size, figure_size)
+ # )
+ font = ImageFont.load_default()
+ width, height = plot.size
+ description, crop_coordinates = self.inverse_build(conditional)
+ for (representation, bbox), color in zip(description, cycle(COLOR_PALETTE)):
+ annotation = self.representation_to_annotation(representation)
+ # class_label = label_for_category_no(annotation.category_id) + ' ' + additional_parameters_string(annotation)
+ class_label = label_for_category_no(annotation.category_id)
+ bbox = absolute_bbox(bbox, width, height)
+ draw.rectangle(bbox, outline=color, width=line_width)
+ draw.text((bbox[0] + line_width, bbox[1] + line_width), class_label, anchor='la', fill=BLACK, font=font)
+ if crop_coordinates is not None:
+ draw.rectangle(absolute_bbox(crop_coordinates, width, height), outline=GRAY_75, width=line_width)
+ return convert_pil_to_tensor(plot) / 127.5 - 1.
diff --git a/lidm/data/conditional_builder/objects_center_points.py b/lidm/data/conditional_builder/objects_center_points.py
new file mode 100644
index 0000000000000000000000000000000000000000..e9b75b85d10ae01b24f7e959fadd9a9b4632023b
--- /dev/null
+++ b/lidm/data/conditional_builder/objects_center_points.py
@@ -0,0 +1,150 @@
+import math
+import random
+import warnings
+from itertools import cycle
+from typing import List, Optional, Tuple, Callable
+
+from PIL import Image as pil_image, ImageDraw as pil_img_draw, ImageFont
+from more_itertools.recipes import grouper
+from .utils import COLOR_PALETTE, WHITE, GRAY_75, BLACK, additional_parameters_string, pad_list, get_circle_size, \
+ get_plot_font_size, absolute_bbox
+from ..helper_types import BoundingBox, Annotation, Image
+from torch import LongTensor, Tensor
+from torchvision.transforms import PILToTensor
+
+
+pil_to_tensor = PILToTensor()
+
+
+def convert_pil_to_tensor(image: Image) -> Tensor:
+ with warnings.catch_warnings():
+ # to filter PyTorch UserWarning as described here: https://github.com/pytorch/vision/issues/2194
+ warnings.simplefilter("ignore")
+ return pil_to_tensor(image)
+
+
+class ObjectsCenterPointsConditionalBuilder:
+ def __init__(self, no_object_classes: int, no_max_objects: int, no_tokens: int, num_beams: int):
+ self.no_object_classes = no_object_classes
+ self.no_max_objects = no_max_objects
+ self.no_tokens = no_tokens
+ # self.no_sections = int(math.sqrt(self.no_tokens))
+ self.no_sections = (self.no_tokens // num_beams, num_beams) # (width, height)
+
+ @property
+ def none(self) -> int:
+ return self.no_tokens - 1
+
+ @property
+ def object_descriptor_length(self) -> int:
+ return 2
+
+ @property
+ def empty_tuple(self) -> Tuple:
+ return (self.none,) * self.object_descriptor_length
+
+ @property
+ def embedding_dim(self) -> int:
+ return self.no_max_objects * self.object_descriptor_length
+
+ def tokenize_coordinates(self, x: float, y: float) -> int:
+ """
+ Express 2d coordinates with one number.
+ Example: assume self.no_tokens = 16, then no_sections = 4:
+ 0 0 0 0
+ 0 0 # 0
+ 0 0 0 0
+ 0 0 0 x
+ Then the # position corresponds to token 6, the x position to token 15.
+ @param x: float in [0, 1]
+ @param y: float in [0, 1]
+ @return: discrete tokenized coordinate
+ """
+ x_discrete = int(round(x * (self.no_sections[0] - 1)))
+ y_discrete = int(round(y * (self.no_sections[1] - 1)))
+ return y_discrete * self.no_sections[0] + x_discrete
+
+ def coordinates_from_token(self, token: int) -> (float, float):
+ x = token % self.no_sections[0]
+ y = token // self.no_sections[0]
+ return x / (self.no_sections[0] - 1), y / (self.no_sections[1] - 1)
+
+ def bbox_from_token_pair(self, token1: int, token2: int) -> BoundingBox:
+ x0, y0 = self.coordinates_from_token(token1)
+ x1, y1 = self.coordinates_from_token(token2)
+ # x2, y2 = self.coordinates_from_token(token3)
+ # x3, y3 = self.coordinates_from_token(token4)
+ return x0, y0, x1, y1
+
+ def token_pair_from_bbox(self, bbox: BoundingBox) -> Tuple:
+ # return self.tokenize_coordinates(bbox[0], bbox[1]), self.tokenize_coordinates(bbox[2], bbox[3]), self.tokenize_coordinates(bbox[4], bbox[5]), self.tokenize_coordinates(bbox[6], bbox[7])
+ return self.tokenize_coordinates(bbox[0], bbox[1]), self.tokenize_coordinates(bbox[4], bbox[5])
+
+ def inverse_build(self, conditional: LongTensor) \
+ -> Tuple[List[Tuple[int, Tuple[float, float]]], Optional[BoundingBox]]:
+ conditional_list = conditional.tolist()
+ table_of_content = grouper(conditional_list, self.object_descriptor_length)
+ assert conditional.shape[0] == self.embedding_dim
+ return [
+ (object_tuple[0], self.coordinates_from_token(object_tuple[1]))
+ for object_tuple in table_of_content if object_tuple[0] != self.none
+ ], None
+
+ def plot(self, conditional: LongTensor, label_for_category_no: Callable[[int], str], figure_size: Tuple[int, int],
+ line_width: int = 3, font_size: Optional[int] = None) -> Tensor:
+ plot = pil_image.new('RGB', figure_size, WHITE)
+ draw = pil_img_draw.Draw(plot)
+ circle_size = get_circle_size(figure_size)
+ # font = ImageFont.truetype('/usr/share/fonts/truetype/lato/Lato-Regular.ttf',
+ # size=get_plot_font_size(font_size, figure_size))
+ font = ImageFont.load_default()
+ width, height = plot.size
+ description, crop_coordinates = self.inverse_build(conditional)
+ for (representation, (x, y)), color in zip(description, cycle(COLOR_PALETTE)):
+ x_abs, y_abs = x * width, y * height
+ ann = self.representation_to_annotation(representation)
+ label = label_for_category_no(ann.category_id) + ' ' + additional_parameters_string(ann)
+ ellipse_bbox = [x_abs - circle_size, y_abs - circle_size, x_abs + circle_size, y_abs + circle_size]
+ draw.ellipse(ellipse_bbox, fill=color, width=0)
+ draw.text((x_abs, y_abs), label, anchor='md', fill=BLACK, font=font)
+ if crop_coordinates is not None:
+ draw.rectangle(absolute_bbox(crop_coordinates, width, height), outline=GRAY_75, width=line_width)
+ return convert_pil_to_tensor(plot) / 127.5 - 1.
+
+ def object_representation(self, annotation: Annotation) -> int:
+ return annotation.category_id
+
+ def representation_to_annotation(self, representation: int) -> Annotation:
+ category_id = representation % self.no_object_classes
+ # noinspection PyTypeChecker
+ return Annotation(
+ bbox=None,
+ category_id=category_id,
+ )
+
+ def _make_object_descriptors(self, annotations: List[Annotation]) -> List[Tuple[int, ...]]:
+ object_tuples = [
+ (self.object_representation(a),
+ self.tokenize_coordinates(a.center[0], a.center[1]))
+ for a in annotations
+ ]
+ empty_tuple = (self.none, self.none)
+ object_tuples = pad_list(object_tuples, empty_tuple, self.no_max_objects)
+ return object_tuples
+
+ def build(self, annotations: List[Annotation]) \
+ -> LongTensor:
+ if len(annotations) == 0:
+ warnings.warn('Did not receive any annotations.')
+
+ random.shuffle(annotations)
+ if len(annotations) > self.no_max_objects:
+ warnings.warn('Received more annotations than allowed.')
+ annotations = annotations[:self.no_max_objects]
+
+ object_tuples = self._make_object_descriptors(annotations)
+ flattened = [token for tuple_ in object_tuples for token in tuple_]
+ assert len(flattened) == self.embedding_dim
+ assert all(0 <= value < self.no_tokens for value in flattened)
+
+ return LongTensor(flattened)
diff --git a/lidm/data/conditional_builder/utils.py b/lidm/data/conditional_builder/utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..e162ccdde4e61b535c85f7a0ca58c9101d184d79
--- /dev/null
+++ b/lidm/data/conditional_builder/utils.py
@@ -0,0 +1,188 @@
+import importlib
+from typing import List, Any, Tuple, Optional
+
+import numpy as np
+from ..helper_types import BoundingBox, Annotation
+
+# source: seaborn, color palette tab10
+COLOR_PALETTE = [(30, 118, 179), (255, 126, 13), (43, 159, 43), (213, 38, 39), (147, 102, 188),
+ (139, 85, 74), (226, 118, 193), (126, 126, 126), (187, 188, 33), (22, 189, 206)]
+BLACK = (0, 0, 0)
+GRAY_75 = (63, 63, 63)
+GRAY_50 = (127, 127, 127)
+GRAY_25 = (191, 191, 191)
+WHITE = (255, 255, 255)
+FULL_CROP = (0., 0., 1., 1.)
+
+
+def corners_3d_to_2d(corners3d):
+ """
+ Args:
+ corners3d: (N, 8, 2)
+ Returns:
+ corners2d: (N, 4, 2)
+ """
+ # select pairs to reorganize
+ mask_0_3 = corners3d[:, 0:4, 0].argmax(1) // 2 != 0
+ mask_4_7 = corners3d[:, 4:8, 0].argmin(1) // 2 != 0
+
+ # reorganize corners in the order of (bottom-right, bottom-left)
+ corners3d[mask_0_3, 0:4] = corners3d[mask_0_3][:, [2, 3, 0, 1]]
+ # reorganize corners in the order of (top-left, top-right)
+ corners3d[mask_4_7, 4:8] = corners3d[mask_4_7][:, [2, 3, 0, 1]]
+
+ # calculate corners in order
+ bot_r = np.stack([corners3d[:, 0:2, 0].max(1), corners3d[:, 0:2, 1].min(1)], axis=-1)
+ bot_l = np.stack([corners3d[:, 2:4, 0].min(1), corners3d[:, 2:4, 1].min(1)], axis=-1)
+ top_l = np.stack([corners3d[:, 4:6, 0].min(1), corners3d[:, 4:6, 1].max(1)], axis=-1)
+ top_r = np.stack([corners3d[:, 6:8, 0].max(1), corners3d[:, 6:8, 1].max(1)], axis=-1)
+
+ return np.stack([bot_r, bot_l, top_l, top_r], axis=1)
+
+
+def rotate_points_along_z(points, angle):
+ """
+ Args:
+ points: (N, 3 + C)
+ angle: angle along z-axis, angle increases x ==> y
+ Returns:
+
+ """
+ cosa = np.cos(angle)
+ sina = np.sin(angle)
+ zeros = np.zeros(points.shape[0])
+ ones = np.ones(points.shape[0])
+ rot_matrix = np.stack((
+ cosa, sina, zeros,
+ -sina, cosa, zeros,
+ zeros, zeros, ones)).reshape((-1, 3, 3))
+ points_rot = np.matmul(points[:, :, 0:3], rot_matrix)
+ points_rot = np.concatenate((points_rot, points[:, :, 3:]), axis=-1)
+ return points_rot
+
+
+def boxes_to_corners_3d(boxes3d):
+ """
+ 7 -------- 4
+ /| /|
+ 6 -------- 5 .
+ | | | |
+ . 3 -------- 0
+ |/ |/
+ 2 -------- 1
+ Args:
+ boxes3d: (N, 7) [x, y, z, dx, dy, dz, heading], (x, y, z) is the box center
+
+ Returns:
+ corners3d: (N, 8, 3)
+ """
+ template = np.array(
+ [[1, 1, -1], [1, -1, -1], [-1, -1, -1], [-1, 1, -1],
+ [1, 1, 1], [1, -1, 1], [-1, -1, 1], [-1, 1, 1]],
+ ) / 2
+
+ # corners3d = boxes3d[:, None, 3:6].repeat(1, 8, 1) * template[None, :, :]
+ corners3d = np.tile(boxes3d[:, None, 3:6], (1, 8, 1)) * template[None, :, :]
+ corners3d = rotate_points_along_z(corners3d.reshape((-1, 8, 3)), boxes3d[:, 6]).reshape((-1, 8, 3))
+ corners3d += boxes3d[:, None, 0:3]
+
+ return corners3d
+
+
+def intersection_area(rectangle1: BoundingBox, rectangle2: BoundingBox) -> float:
+ """
+ Give intersection area of two rectangles.
+ @param rectangle1: (x0, y0, w, h) of first rectangle
+ @param rectangle2: (x0, y0, w, h) of second rectangle
+ """
+ rectangle1 = rectangle1[0], rectangle1[1], rectangle1[0] + rectangle1[2], rectangle1[1] + rectangle1[3]
+ rectangle2 = rectangle2[0], rectangle2[1], rectangle2[0] + rectangle2[2], rectangle2[1] + rectangle2[3]
+ x_overlap = max(0., min(rectangle1[2], rectangle2[2]) - max(rectangle1[0], rectangle2[0]))
+ y_overlap = max(0., min(rectangle1[3], rectangle2[3]) - max(rectangle1[1], rectangle2[1]))
+ return x_overlap * y_overlap
+
+
+def horizontally_flip_bbox(bbox: BoundingBox) -> BoundingBox:
+ return 1 - (bbox[0] + bbox[2]), bbox[1], bbox[2], bbox[3]
+
+
+def absolute_bbox(relative_bbox: BoundingBox, width: int, height: int) -> Tuple[int, int, int, int]:
+ bbox = relative_bbox
+ # bbox = bbox[0] * width, bbox[1] * height, (bbox[0] + bbox[2]) * width, (bbox[1] + bbox[3]) * height
+ bbox = bbox[0] * width, bbox[1] * height, bbox[2] * width, bbox[3] * height
+ # return int(bbox[0]), int(bbox[1]), int(bbox[2]), int(bbox[3])
+ x1, x2 = min(int(bbox[2]), int(bbox[0])), max(int(bbox[2]), int(bbox[0]))
+ y1, y2 = min(int(bbox[3]), int(bbox[1])), max(int(bbox[3]), int(bbox[1]))
+ if x1 == x2:
+ x2 += 1
+ if y1 == y2:
+ y2 += 1
+ return x1, y1, x2, y2
+
+
+def pad_list(list_: List, pad_element: Any, pad_to_length: int) -> List:
+ return list_ + [pad_element for _ in range(pad_to_length - len(list_))]
+
+
+def rescale_annotations(annotations: List[Annotation], crop_coordinates: BoundingBox, flip: bool) -> \
+ List[Annotation]:
+ def clamp(x: float):
+ return max(min(x, 1.), 0.)
+
+ def rescale_bbox(bbox: BoundingBox) -> BoundingBox:
+ x0 = clamp((bbox[0] - crop_coordinates[0]) / crop_coordinates[2])
+ y0 = clamp((bbox[1] - crop_coordinates[1]) / crop_coordinates[3])
+ w = min(bbox[2] / crop_coordinates[2], 1 - x0)
+ h = min(bbox[3] / crop_coordinates[3], 1 - y0)
+ if flip:
+ x0 = 1 - (x0 + w)
+ return x0, y0, w, h
+
+ return [a._replace(bbox=rescale_bbox(a.bbox)) for a in annotations]
+
+
+def filter_annotations(annotations: List[Annotation], crop_coordinates: BoundingBox) -> List:
+ return [a for a in annotations if intersection_area(a.bbox, crop_coordinates) > 0.0]
+
+
+def additional_parameters_string(annotation: Annotation, short: bool = True) -> str:
+ sl = slice(1) if short else slice(None)
+ string = ''
+ if not (annotation.is_group_of or annotation.is_occluded or annotation.is_depiction or annotation.is_inside):
+ return string
+ if annotation.is_group_of:
+ string += 'group'[sl] + ','
+ if annotation.is_occluded:
+ string += 'occluded'[sl] + ','
+ if annotation.is_depiction:
+ string += 'depiction'[sl] + ','
+ if annotation.is_inside:
+ string += 'inside'[sl]
+ return '(' + string.strip(",") + ')'
+
+
+def get_plot_font_size(font_size: Optional[int], figure_size: Tuple[int, int]) -> int:
+ if font_size is None:
+ font_size = 10
+ if max(figure_size) >= 256:
+ font_size = 12
+ if max(figure_size) >= 512:
+ font_size = 15
+ return font_size
+
+
+def get_circle_size(figure_size: Tuple[int, int]) -> int:
+ circle_size = 2
+ if max(figure_size) >= 256:
+ circle_size = 3
+ if max(figure_size) >= 512:
+ circle_size = 4
+ return circle_size
+
+
+def load_object_from_string(object_string: str) -> Any:
+ """
+ Source: https://stackoverflow.com/a/10773699
+ """
+ module_name, class_name = object_string.rsplit(".", 1)
+ return getattr(importlib.import_module(module_name), class_name)
diff --git a/lidm/data/helper_types.py b/lidm/data/helper_types.py
new file mode 100644
index 0000000000000000000000000000000000000000..ae525c7775ad55090511ef021db5867ce030e421
--- /dev/null
+++ b/lidm/data/helper_types.py
@@ -0,0 +1,20 @@
+from typing import Tuple, Optional, NamedTuple, Union, List
+from PIL.Image import Image as pil_image
+from torch import Tensor
+
+try:
+ from typing import Literal
+except ImportError:
+ from typing_extensions import Literal
+
+Image = Union[Tensor, pil_image]
+# BoundingBox = Tuple[float, float, float, float] # x0, y0, w, h | x0, y0, x1, y1
+# BoundingBox3D = Tuple[float, float, float, float, float, float] # x0, y0, z0, l, w, h
+BoundingBox = Tuple[float, float, float, float] # corner coordinates (x,y) in the order of bottom-right -> bottom-left -> top-left -> top-right
+Center = Tuple[float, float]
+
+
+class Annotation(NamedTuple):
+ category_id: int
+ bbox: Optional[BoundingBox] = None
+ center: Optional[Center] = None
diff --git a/lidm/data/kitti.py b/lidm/data/kitti.py
new file mode 100644
index 0000000000000000000000000000000000000000..103426449bbda265d817ea2641dd48c3194d2c4e
--- /dev/null
+++ b/lidm/data/kitti.py
@@ -0,0 +1,345 @@
+import glob
+import os
+import pickle
+import numpy as np
+import yaml
+from PIL import Image
+import xml.etree.ElementTree as ET
+
+from lidm.data.base import DatasetBase
+from .annotated_dataset import Annotated3DObjectsDataset
+from .conditional_builder.utils import corners_3d_to_2d
+from .helper_types import Annotation
+from ..utils.lidar_utils import pcd2range, pcd2coord2d, range2pcd
+
+# TODO add annotation categories and semantic categories
+CATEGORIES = ['ignore', 'car', 'bicycle', 'motorcycle', 'truck', 'other-vehicle', 'person', 'bicyclist', 'motorcyclist',
+ 'road', 'parking', 'sidewalk', 'other-ground', 'building', 'fence', 'vegetation', 'trunk', 'terrain',
+ 'pole', 'traffic-sign']
+CATE2LABEL = {k: v for v, k in enumerate(CATEGORIES)} # 0: invalid, 1~10: categories
+LABEL2RGB = np.array([(0, 0, 0), (0, 0, 142), (119, 11, 32), (0, 0, 230), (0, 0, 70), (0, 0, 90), (220, 20, 60),
+ (255, 0, 0), (0, 0, 110), (128, 64, 128), (250, 170, 160), (244, 35, 232), (230, 150, 140),
+ (70, 70, 70), (190, 153, 153), (107, 142, 35), (0, 80, 100), (230, 150, 140), (153, 153, 153),
+ (220, 220, 0)])
+CAMERAS = ['CAM_FRONT']
+BBOX_CATS = ['car', 'people', 'cycle']
+BBOX_CAT2LABEL = {'car': 0, 'truck': 0, 'bus': 0, 'caravan': 0, 'person': 1, 'rider': 2, 'motorcycle': 2, 'bicycle': 2}
+
+# train + test
+SEM_KITTI_TRAIN_SET = ['00', '01', '02', '03', '04', '05', '06', '07', '09', '10']
+KITTI_TRAIN_SET = SEM_KITTI_TRAIN_SET + ['11', '12', '13', '14', '15', '16', '17', '18', '19', '20', '21']
+KITTI360_TRAIN_SET = ['00', '02', '04', '05', '06', '07', '09', '10'] + ['08'] # partial test data at '02' sequence
+CAM_KITTI360_TRAIN_SET = ['00', '04', '05', '06', '07', '08', '09', '10'] # cam mismatch lidar in '02'
+
+# validation
+SEM_KITTI_VAL_SET = KITTI_VAL_SET = ['08']
+CAM_KITTI360_VAL_SET = KITTI360_VAL_SET = ['03']
+
+
+class KITTIBase(DatasetBase):
+ def __init__(self, **kwargs):
+ super().__init__(**kwargs)
+ self.dataset_name = 'kitti'
+ self.num_sem_cats = kwargs['dataset_config'].num_sem_cats + 1
+
+ @staticmethod
+ def load_lidar_sweep(path):
+ scan = np.fromfile(path, dtype=np.float32)
+ scan = scan.reshape((-1, 4))
+ points = scan[:, 0:3] # get xyz
+ return points
+
+ def load_semantic_map(self, path, pcd):
+ raise NotImplementedError
+
+ def load_camera(self, path):
+ raise NotImplementedError
+
+ def __getitem__(self, idx):
+ example = dict()
+ data_path = self.data[idx]
+ # lidar point cloud
+ sweep = self.load_lidar_sweep(data_path)
+
+ if self.lidar_transform:
+ sweep, _ = self.lidar_transform(sweep, None)
+
+ if self.condition_key == 'segmentation':
+ # semantic maps
+ proj_range, sem_map = self.load_semantic_map(data_path, sweep)
+ example[self.condition_key] = sem_map
+ else:
+ proj_range, _ = pcd2range(sweep, self.img_size, self.fov, self.depth_range)
+ proj_range, proj_mask = self.process_scan(proj_range)
+ example['image'], example['mask'] = proj_range, proj_mask
+ if self.return_pcd:
+ reproj_sweep, _, _ = range2pcd(proj_range[0] * .5 + .5, self.fov, self.depth_range, self.depth_scale, self.log_scale)
+ example['raw'] = sweep
+ example['reproj'] = reproj_sweep.astype(np.float32)
+
+ # image degradation
+ if self.degradation_transform:
+ degraded_proj_range = self.degradation_transform(proj_range)
+ example['degraded_image'] = degraded_proj_range
+
+ # cameras
+ if self.condition_key == 'camera':
+ cameras = self.load_camera(data_path)
+ example[self.condition_key] = cameras
+
+ return example
+
+
+class SemanticKITTIBase(KITTIBase):
+ def __init__(self, **kwargs):
+ super().__init__(**kwargs)
+ assert self.condition_key in ['segmentation'] # for segmentation input only
+ self.label2rgb = LABEL2RGB
+
+ def prepare_data(self):
+ # read data paths from KITTI
+ for seq_id in eval('SEM_KITTI_%s_SET' % self.split.upper()):
+ self.data.extend(glob.glob(os.path.join(
+ self.data_root, f'dataset/sequences/{seq_id}/velodyne/*.bin')))
+ # read label mapping
+ data_config = yaml.safe_load(open('./data/config/semantic-kitti.yaml', 'r'))
+ remap_dict = data_config["learning_map"]
+ max_key = max(remap_dict.keys())
+ self.learning_map = np.zeros((max_key + 100), dtype=np.int32)
+ self.learning_map[list(remap_dict.keys())] = list(remap_dict.values())
+
+ def load_semantic_map(self, path, pcd):
+ label_path = path.replace('velodyne', 'labels').replace('.bin', '.label')
+ labels = np.fromfile(label_path, dtype=np.uint32)
+ labels = labels.reshape((-1))
+ labels = labels & 0xFFFF # semantic label in lower half
+ labels = self.learning_map[labels]
+
+ proj_range, sem_map = pcd2range(pcd, self.img_size, self.fov, self.depth_range, labels=labels)
+ # sem_map = np.expand_dims(sem_map, axis=0).astype(np.int64)
+ sem_map = sem_map.astype(np.int64)
+ if self.filtered_map_cats is not None:
+ sem_map[np.isin(sem_map, self.filtered_map_cats)] = 0 # set filtered category as noise
+ onehot = np.eye(self.num_sem_cats, dtype=np.float32)[sem_map].transpose(2, 0, 1)
+ return proj_range, onehot
+
+
+class SemanticKITTITrain(SemanticKITTIBase):
+ def __init__(self, **kwargs):
+ super().__init__(data_root='./dataset/SemanticKITTI', split='train', **kwargs)
+
+
+class SemanticKITTIValidation(SemanticKITTIBase):
+ def __init__(self, **kwargs):
+ super().__init__(data_root='./dataset/SemanticKITTI', split='val', **kwargs)
+
+
+class KITTI360Base(KITTIBase):
+ def __init__(self, split_per_view=None, **kwargs):
+ super().__init__(**kwargs)
+ self.split_per_view = split_per_view
+ if self.condition_key == 'camera':
+ assert self.split_per_view is not None, 'For camera-to-lidar, need to specify split_per_view'
+
+ def prepare_data(self):
+ # read data paths
+ self.data = []
+ if self.condition_key == 'camera':
+ seq_list = eval('CAM_KITTI360_%s_SET' % self.split.upper())
+ else:
+ seq_list = eval('KITTI360_%s_SET' % self.split.upper())
+ for seq_id in seq_list:
+ self.data.extend(glob.glob(os.path.join(
+ self.data_root, f'data_3d_raw/2013_05_28_drive_00{seq_id}_sync/velodyne_points/data/*.bin')))
+
+ def random_drop_camera(self, camera_list):
+ if np.random.rand() < self.aug_config['camera_drop'] and self.split == 'train':
+ camera_list = [np.zeros_like(c) if i != len(camera_list) // 2 else c for i, c in enumerate(camera_list)] # keep the middle view only
+ return camera_list
+
+ def load_camera(self, path):
+ camera_path = path.replace('data_3d_raw', 'data_2d_camera').replace('velodyne_points/data', 'image_00/data_rect').replace('.bin', '.png')
+ camera = np.array(Image.open(camera_path)).astype(np.float32) / 255.
+ camera = camera.transpose(2, 0, 1)
+ if self.view_transform:
+ camera = self.view_transform(camera)
+ camera_list = np.split(camera, self.split_per_view, axis=2) # split into n chunks as different views
+ camera_list = self.random_drop_camera(camera_list)
+ return camera_list
+
+
+class KITTI360Train(KITTI360Base):
+ def __init__(self, **kwargs):
+ super().__init__(data_root='./dataset/KITTI-360', split='train', **kwargs)
+
+
+class KITTI360Validation(KITTI360Base):
+ def __init__(self, **kwargs):
+ super().__init__(data_root='./dataset/KITTI-360', split='val', **kwargs)
+
+
+class AnnotatedKITTI360Base(Annotated3DObjectsDataset, KITTI360Base):
+ def __init__(self, **kwargs):
+ self.id_bbox_dict = dict()
+ self.id_label_dict = dict()
+
+ Annotated3DObjectsDataset.__init__(self, **kwargs)
+ KITTI360Base.__init__(self, **kwargs)
+ assert self.condition_key in ['center', 'bbox'] # for annotated images only
+
+ @staticmethod
+ def parseOpencvMatrix(node):
+ rows = int(node.find('rows').text)
+ cols = int(node.find('cols').text)
+ data = node.find('data').text.split(' ')
+
+ mat = []
+ for d in data:
+ d = d.replace('\n', '')
+ if len(d) < 1:
+ continue
+ mat.append(float(d))
+ mat = np.reshape(mat, [rows, cols])
+ return mat
+
+ def parseVertices(self, child):
+ transform = self.parseOpencvMatrix(child.find('transform'))
+ R = transform[:3, :3]
+ T = transform[:3, 3]
+ vertices = self.parseOpencvMatrix(child.find('vertices'))
+ vertices = np.matmul(R, vertices.transpose()).transpose() + T
+ return vertices
+
+ def parse_bbox_xml(self, path):
+ tree = ET.parse(path)
+ root = tree.getroot()
+
+ bbox_dict = dict()
+ label_dict = dict()
+ for child in root:
+ if child.find('transform') is None:
+ continue
+
+ label_name = child.find('label').text
+ if label_name not in BBOX_CAT2LABEL:
+ continue
+
+ label = BBOX_CAT2LABEL[label_name]
+ timestamp = int(child.find('timestamp').text)
+ # verts = self.parseVertices(child)
+ verts = self.parseOpencvMatrix(child.find('vertices'))[:8]
+ if timestamp in bbox_dict:
+ bbox_dict[timestamp].append(verts)
+ label_dict[timestamp].append(label)
+ else:
+ bbox_dict[timestamp] = [verts]
+ label_dict[timestamp] = [label]
+ return bbox_dict, label_dict
+
+ def prepare_data(self):
+ KITTI360Base.prepare_data(self)
+
+ self.data = [p for p in self.data if '2013_05_28_drive_0008_sync' not in p] # remove unlabeled sequence 08
+ seq_list = eval('KITTI360_%s_SET' % self.split.upper())
+ for seq_id in seq_list:
+ if seq_id != '08':
+ xml_path = os.path.join(self.data_root, f'data_3d_bboxes/train/2013_05_28_drive_00{seq_id}_sync.xml')
+ bbox_dict, label_dict = self.parse_bbox_xml(xml_path)
+ self.id_bbox_dict[seq_id] = bbox_dict
+ self.id_label_dict[seq_id] = label_dict
+
+ def load_annotation(self, path):
+ seq_id = path.split('/')[-4].split('_')[-2][-2:]
+ timestamp = int(path.split('/')[-1].replace('.bin', ''))
+ verts_list = self.id_bbox_dict[seq_id][timestamp]
+ label_list = self.id_label_dict[seq_id][timestamp]
+
+ if self.condition_key == 'bbox':
+ points = np.stack(verts_list)
+ elif self.condition_key == 'center':
+ points = (verts_list[0] + verts_list[6]) / 2.
+ else:
+ raise NotImplementedError
+ labels = np.array([label_list])
+ if self.anno_transform:
+ points, labels = self.anno_transform(points, labels)
+ return points, labels
+
+ def __getitem__(self, idx):
+ example = dict()
+ data_path = self.data[idx]
+
+ # lidar point cloud
+ sweep = self.load_lidar_sweep(data_path)
+
+ # annotations
+ bbox_points, bbox_labels = self.load_annotation(data_path)
+
+ if self.lidar_transform:
+ sweep, bbox_points = self.lidar_transform(sweep, bbox_points)
+
+ # point cloud -> range
+ proj_range, _ = pcd2range(sweep, self.img_size, self.fov, self.depth_range)
+ proj_range, proj_mask = self.process_scan(proj_range)
+ example['image'], example['mask'] = proj_range, proj_mask
+ if self.return_pcd:
+ example['reproj'] = sweep
+
+ # annotation -> range
+ # NOTE: do not need to transform bbox points along with lidar, since their coordinates are based on range-image space instead of 3D space
+ proj_bbox_points, proj_bbox_labels = pcd2coord2d(bbox_points, self.fov, self.depth_range, labels=bbox_labels)
+ builder = self.conditional_builders[self.condition_key]
+ if self.condition_key == 'bbox':
+ proj_bbox_points = corners_3d_to_2d(proj_bbox_points)
+ annotations = [Annotation(bbox=bbox.flatten(), category_id=label) for bbox, label in
+ zip(proj_bbox_points, proj_bbox_labels)]
+ else:
+ annotations = [Annotation(center=center, category_id=label) for center, label in
+ zip(proj_bbox_points, proj_bbox_labels)]
+ example[self.condition_key] = builder.build(annotations)
+
+ return example
+
+
+class AnnotatedKITTI360Train(AnnotatedKITTI360Base):
+ def __init__(self, **kwargs):
+ super().__init__(data_root='./dataset/KITTI-360', split='train', cats=BBOX_CATS, **kwargs)
+
+
+class AnnotatedKITTI360Validation(AnnotatedKITTI360Base):
+ def __init__(self, **kwargs):
+ super().__init__(data_root='./dataset/KITTI-360', split='train', cats=BBOX_CATS, **kwargs)
+
+
+class KITTIImageBase(KITTIBase):
+ """
+ Range ImageSet only combining KITTI-360 and SemanticKITTI
+
+ #Samples (Training): 98014, #Samples (Val): 3511
+
+ """
+ def __init__(self, **kwargs):
+ super().__init__(**kwargs)
+ assert self.condition_key in [None, 'image'] # for image input only
+
+ def prepare_data(self):
+ # read data paths from KITTI-360
+ self.data = []
+ for seq_id in eval('KITTI360_%s_SET' % self.split.upper()):
+ self.data.extend(glob.glob(os.path.join(
+ self.data_root, f'KITTI-360/data_3d_raw/2013_05_28_drive_00{seq_id}_sync/velodyne_points/data/*.bin')))
+
+ # read data paths from KITTI
+ for seq_id in eval('KITTI_%s_SET' % self.split.upper()):
+ self.data.extend(glob.glob(os.path.join(
+ self.data_root, f'SemanticKITTI/dataset/sequences/{seq_id}/velodyne/*.bin')))
+
+
+class KITTIImageTrain(KITTIImageBase):
+ def __init__(self, **kwargs):
+ super().__init__(data_root='./dataset', split='train', **kwargs)
+
+
+class KITTIImageValidation(KITTIImageBase):
+ def __init__(self, **kwargs):
+ super().__init__(data_root='./dataset', split='val', **kwargs)
diff --git a/lidm/eval/README.md b/lidm/eval/README.md
new file mode 100644
index 0000000000000000000000000000000000000000..33e8f8998cf905ac22e4b64a1b4ba87153ca880e
--- /dev/null
+++ b/lidm/eval/README.md
@@ -0,0 +1,95 @@
+# Evaluation Toolbox for LiDAR Generation
+
+This directory is a **self-contained**, **memory-friendly** and mostly **CUDA-accelerated** toolbox of multiple evaluation metrics for LiDAR generative models, including:
+* Perceptual metrics (our proposed):
+ * Fréchet Range Image Distance (**FRID**)
+ * Fréchet Sparse Volume Distance (**FSVD**)
+ * Fréchet Point-based Volume Distance (**FPVD**)
+* Statistical metrics (proposed in [Learning Representations and Generative Models for 3D Point Clouds](https://arxiv.org/abs/1707.02392)):
+ * Minimum Matching Distance (**MMD**)
+ * Jensen-Shannon Divergence (**JSD**)
+* Statistical pairwise metrics (for reconstruction only):
+ * Chamfer Distance (**CD**)
+ * Earth Mover's Distance (**EMD**)
+
+## Citation
+
+If you find this project useful in your research, please consider citing:
+```
+@article{ran2024towards,
+ title={Towards Realistic Scene Generation with LiDAR Diffusion Models},
+ author={Ran, Haoxi and Guizilini, Vitor and Wang, Yue},
+ journal={arXiv preprint arXiv:2404.00815},
+ year={2024}
+}
+```
+
+
+## Dependencies
+
+### Basic (install through **pip**):
+* scipy
+* numpy
+* torch
+* pyyaml
+
+### Required by FSVD and FPVD:
+* [Torchsparse v1.4.0](https://github.com/mit-han-lab/torchsparse/tree/v1.4.0) (pip install git+https://github.com/mit-han-lab/torchsparse.git@v1.4.0)
+* [Google Sparse Hash library](https://github.com/sparsehash/sparsehash) (apt-get install libsparsehash-dev **or** compile locally and update variable CPLUS_INCLUDE_PATH with directory path)
+
+
+## Model Zoo
+
+To evaluate with perceptual metrics on different types of LiDAR data, you can download all models through:
+* this [google drive link](https://drive.google.com/file/d/1Ml4p4_nMlwLkSp7JB528GJv2_HxO8v1i/view?usp=drive_link) in the .zip file
+
+or
+* the **full directory** of one specific model:
+
+### 64-beam LiDAR (trained on [SemanticKITTI](http://semantic-kitti.org/dataset.html)):
+
+| Metric | Model | Arch | Link | Code | Comments |
+|:------:|:-------------------------------------------------------------------------------------------:|:-----------------------:|:-------------------------------------------------------------------------------------------------------:|:-----------------------------------------------------------------|---------------------------------------------------------------------------|
+| FRID | [RangeNet++](https://www.ipb.uni-bonn.de/wp-content/papercite-data/pdf/milioto2019iros.pdf) | DarkNet21-based UNet | [Google Drive](https://drive.google.com/drive/folders/1ZS8KOoxB9hjB6kwKbH5Zfc8O5qJlKsbl?usp=drive_link) | [./models/rangenet/model.py](./models/rangenet/model.py) | range image input (our trained model without the need of remission input) |
+| FSVD | [MinkowskiNet](https://arxiv.org/abs/1904.08755) | Sparse UNet | [Google Drive](https://drive.google.com/drive/folders/1zN12ZEvjIvo4PCjAsncgC22yvtRrCCMe?usp=drive_link) | [./models/minkowskinet/model.py](./models/minkowskinet/model.py) | point cloud input |
+| FPVD | [SPVCNN](https://arxiv.org/abs/2007.16100) | Point-Voxel Sparse UNet | [Google Drive](https://drive.google.com/drive/folders/1oEm3qpxfGetiVAfXIvecawEiFqW79M6B?usp=drive_link) | [./models/spvcnn/model.py](./models/spvcnn/model.py) | point cloud input |
+
+
+### 32-beam LiDAR (trained on [nuScenes](https://www.nuscenes.org/nuscenes)):
+
+| Metric | Model | Arch | Link | Code | Comments |
+|:------:|:------------------------------------------------:|:-----------------------:|:-------------------------------------------------------------------------------------------------------:|:-----------------------------------------------------------------|-------------------|
+| FSVD | [MinkowskiNet](https://arxiv.org/abs/1904.08755) | Sparse UNet | [Google Drive](https://drive.google.com/drive/folders/1oZIS9FlklCQ6dlh3TZ8Junir7QwgT-Me?usp=drive_link) | [./models/minkowskinet/model.py](./models/minkowskinet/model.py) | point cloud input |
+| FPVD | [SPVCNN](https://arxiv.org/abs/2007.16100) | Point-Voxel Sparse UNet | [Google Drive](https://drive.google.com/drive/folders/1F69RbprAoT6MOJ7iI0KHjxuq-tbeqGiR?usp=drive_link) | [./models/spvcnn/model.py](./models/spvcnn/model.py) | point cloud input |
+
+
+## Usage
+
+1. Place the unzipped `pretrained_weights` folder under the root python directory **or** modify the `DEFAULT_ROOT` variable in the `__init__.py`.
+2. Prepare input data, including the synthesized samples and the reference dataset. **Note**: The reference data should be the **point clouds projected back from range images** instead of raw point clouds.
+3. Specify the data type (`32` or `64`) and the metrics to evaluate. Options: `mmd`, `jsd`, `frid`, `fsvd`, `fpvd`, `cd`, `emd`.
+4. (Optional) If you want to compute `frid`, `fsvd` or `fpvd` metric, adjust the corresponding batch size through the `MODAL2BATCHSIZE` in file `__init__.py` according to your max GPU memory (default: ~24GB).
+5. Start evaluation and all results will print out!
+
+### Example:
+
+```
+from .eval_utils import evaluate
+
+data = '64' # specify data type to evaluate
+metrics = ['mmd', 'jsd', 'frid', 'fsvd', 'fpvd'] # specify metrics to evaluate
+
+# list of np.float32 array
+# shape of each array: (#points, #dim=3), #dim: xyz coordinate (NOTE: no need to input remission)
+reference = ...
+samples = ...
+
+evaluate(reference, samples, metrics, data)
+```
+
+
+## Acknowledgement
+
+- The implementation of MinkowskiNet and SPVCNN is borrowed from [2DPASS](https://github.com/yanx27/2DPASS).
+- The implementation of RangeNet++ is borrowed from [the official RangeNet++ codebase](https://github.com/PRBonn/lidar-bonnetal).
+- The implementation of Chamfer Distance is adapted from [CD Pytorch Implementation](https://github.com/ThibaultGROUEIX/ChamferDistancePytorch) and Earth Mover's Distance from [MSN official repo](https://github.com/Colin97/MSN-Point-Cloud-Completion).
diff --git a/lidm/eval/__init__.py b/lidm/eval/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..b66df061c883b417576c2c1f6ae26f991c6c64c4
--- /dev/null
+++ b/lidm/eval/__init__.py
@@ -0,0 +1,62 @@
+"""
+@Author: Haoxi Ran
+@Date: 01/03/2024
+@Citation: Towards Realistic Scene Generation with LiDAR Diffusion Models
+
+"""
+
+import os
+
+import torch
+import yaml
+
+from lidm.utils.misc_utils import dict2namespace
+from ..modules.rangenet.model import Model as rangenet
+
+try:
+ from ..modules.spvcnn.model import Model as spvcnn
+ from ..modules.minkowskinet.model import Model as minkowskinet
+except:
+ print('To install torchsparse 1.4.0, please refer to https://github.com/mit-han-lab/torchsparse/tree/74099d10a51c71c14318bce63d6421f698b24f24')
+
+# user settings
+DEFAULT_ROOT = './pretrained_weights'
+MODAL2BATCHSIZE = {'range': 100, 'voxel': 50, 'point_voxel': 25}
+OUTPUT_TEMPLATE = 50 * '-' + '\n|' + 16 * ' ' + '{}:{:.4E}' + 17 * ' ' + '|\n' + 50 * '-'
+
+# eval settings (do not modify)
+VOXEL_SIZE = 0.05
+NUM_SECTORS = 16
+AGG_TYPE = 'depth'
+TYPE2DATASET = {'32': 'nuscenes', '64': 'kitti'}
+DATA_CONFIG = {'64': {'x': [-50, 50], 'y': [-50, 50], 'z': [-3, 1]},
+ '32': {'x': [-30, 30], 'y': [-30, 30], 'z': [-3, 6]}}
+MODALITY2MODEL = {'range': 'rangenet', 'voxel': 'minkowskinet', 'point_voxel': 'spvcnn'}
+DATASET_CONFIG = {'kitti': {'size': [64, 1024], 'fov': [3, -25], 'depth_range': [1.0, 56.0], 'depth_scale': 6},
+ 'nuscenes': {'size': [32, 1024], 'fov': [10, -30], 'depth_range': [1.0, 45.0]}}
+
+
+def build_model(dataset_name, model_name, device='cpu'):
+ # config
+ model_folder = os.path.join(DEFAULT_ROOT, dataset_name, model_name)
+
+ if not os.path.isdir(model_folder):
+ raise Exception('Not Available Pretrained Weights!')
+
+ config = yaml.safe_load(open(os.path.join(model_folder, 'config.yaml'), 'r'))
+ if model_name != 'rangenet':
+ config = dict2namespace(config)
+
+ # build model
+ model = eval(model_name)(config)
+
+ # load checkpoint
+ if model_name == 'rangenet':
+ model.load_pretrained_weights(model_folder)
+ else:
+ ckpt = torch.load(os.path.join(model_folder, 'model.ckpt'), map_location="cpu")
+ model.load_state_dict(ckpt['state_dict'], strict=False)
+ model.to(device)
+ model.eval()
+
+ return model
diff --git a/lidm/eval/compile.sh b/lidm/eval/compile.sh
new file mode 100644
index 0000000000000000000000000000000000000000..4d805f490ef9044309635b31b663c45bceaa4233
--- /dev/null
+++ b/lidm/eval/compile.sh
@@ -0,0 +1,9 @@
+#!/bin/sh
+
+cd modules/chamfer
+python setup.py build_ext --inplace
+
+cd ../emd
+python setup.py build_ext --inplace
+
+cd ..
diff --git a/lidm/eval/eval_utils.py b/lidm/eval/eval_utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..bdfef3d58d5462d135d5f4518848530d826a2cad
--- /dev/null
+++ b/lidm/eval/eval_utils.py
@@ -0,0 +1,138 @@
+"""
+@Author: Haoxi Ran
+@Date: 01/03/2024
+@Citation: Towards Realistic Scene Generation with LiDAR Diffusion Models
+
+"""
+import multiprocessing
+from functools import partial
+
+import numpy as np
+from scipy.spatial.distance import jensenshannon
+from tqdm import tqdm
+
+from . import OUTPUT_TEMPLATE
+from .metric_utils import compute_logits, compute_pairwise_cd, \
+ compute_pairwise_emd, pcd2bev_sum, compute_pairwise_cd_batch, pcd2bev_bin
+from .fid_score import calculate_frechet_distance
+
+
+def evaluate(reference, samples, metrics, data):
+ # perceptual
+ if 'frid' in metrics:
+ compute_frid(reference, samples, data)
+ if 'fsvd' in metrics:
+ compute_fsvd(reference, samples, data)
+ if 'fpvd' in metrics:
+ compute_fpvd(reference, samples, data)
+
+ # reconstruction
+ if 'cd' in metrics:
+ compute_cd(reference, samples)
+ if 'emd' in metrics:
+ compute_emd(reference, samples)
+
+ # statistical
+ if 'jsd' in metrics:
+ compute_jsd(reference, samples, data)
+ if 'mmd' in metrics:
+ compute_mmd(reference, samples, data)
+
+
+def compute_cd(reference, samples):
+ """
+ Calculate score of Chamfer Distance (CD)
+
+ """
+ print('Evaluating (CD) ...')
+ results = []
+ for x, y in zip(reference, samples):
+ d = compute_pairwise_cd(x, y)
+ results.append(d)
+ score = sum(results) / len(results)
+ print(OUTPUT_TEMPLATE.format('CD ', score))
+
+
+def compute_emd(reference, samples):
+ """
+ Calculate score of Earth Mover's Distance (EMD)
+
+ """
+ print('Evaluating (EMD) ...')
+ results = []
+ for x, y in zip(reference, samples):
+ d = compute_pairwise_emd(x, y)
+ results.append(d)
+ score = sum(results) / len(results)
+ print(OUTPUT_TEMPLATE.format('EMD ', score))
+
+
+def compute_mmd(reference, samples, data, dist='cd', verbose=True):
+ """
+ Calculate the score of Minimum Matching Distance (MMD)
+
+ """
+ print('Evaluating (MMD) ...')
+ assert dist in ['cd', 'emd']
+ reference, samples = pcd2bev_bin(data, reference, samples)
+ compute_dist_func = compute_pairwise_cd_batch if dist == 'cd' else compute_pairwise_emd
+ results = []
+ for r in tqdm(reference, disable=not verbose):
+ dists = compute_dist_func(r, samples)
+ results.append(min(dists))
+ score = sum(results) / len(results)
+ print(OUTPUT_TEMPLATE.format('MMD ', score))
+
+
+def compute_jsd(reference, samples, data):
+ """
+ Calculate the score of Jensen-Shannon Divergence (JSD)
+
+ """
+ print('Evaluating (JSD) ...')
+ reference, samples = pcd2bev_sum(data, reference, samples)
+ reference = (reference / np.sum(reference)).flatten()
+ samples = (samples / np.sum(samples)).flatten()
+ score = jensenshannon(reference, samples)
+ print(OUTPUT_TEMPLATE.format('JSD ', score))
+
+
+def compute_fd(reference, samples):
+ mu1, mu2 = np.mean(reference, axis=0), np.mean(samples, axis=0)
+ sigma1, sigma2 = np.cov(reference, rowvar=False), np.cov(samples, rowvar=False)
+ distance = calculate_frechet_distance(mu1, sigma1, mu2, sigma2)
+ return distance
+
+
+def compute_frid(reference, samples, data):
+ """
+ Calculate the score of Fréchet Range Image Distance (FRID)
+
+ """
+ print('Evaluating (FRID) ...')
+ gt_logits, samples_logits = compute_logits(data, 'range', reference, samples)
+ score = compute_fd(gt_logits, samples_logits)
+ print(OUTPUT_TEMPLATE.format('FRID', score))
+
+
+def compute_fsvd(reference, samples, data):
+ """
+ Calculate the score of Fréchet Sparse Volume Distance (FSVD)
+
+ """
+ print('Evaluating (FSVD) ...')
+ gt_logits, samples_logits = compute_logits(data, 'voxel', reference, samples)
+ score = compute_fd(gt_logits, samples_logits)
+ print(OUTPUT_TEMPLATE.format('FSVD', score))
+
+
+def compute_fpvd(reference, samples, data):
+ """
+ Calculate the score of Fréchet Point-based Volume Distance (FPVD)
+
+ """
+ print('Evaluating (FPVD) ...')
+ gt_logits, samples_logits = compute_logits(data, 'point_voxel', reference, samples)
+ score = compute_fd(gt_logits, samples_logits)
+ print(OUTPUT_TEMPLATE.format('FPVD', score))
+
diff --git a/lidm/eval/fid_score.py b/lidm/eval/fid_score.py
new file mode 100644
index 0000000000000000000000000000000000000000..56fbb41329017636216cb6458b9cc82440152986
--- /dev/null
+++ b/lidm/eval/fid_score.py
@@ -0,0 +1,191 @@
+"""Calculates the Frechet Inception Distance (FID) to evalulate GANs
+The FID metric calculates the distance between two distributions of images.
+Typically, we have summary statistics (mean & covariance matrix) of one
+of these distributions, while the 2nd distribution is given by a GAN.
+When run as a stand-alone program, it compares the distribution of
+images that are stored as PNG/JPEG at a specified location with a
+distribution given by summary statistics (in pickle format).
+The FID is calculated by assuming that X_1 and X_2 are the activations of
+the pool_3 layer of the inception net for generated samples and real world
+samples respectively.
+See --help to see further details.
+Code adapted from https://github.com/bioinf-jku/TTUR to use PyTorch instead
+of Tensorflow
+Copyright 2018 Institute of Bioinformatics, JKU Linz
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+ http://www.apache.org/licenses/LICENSE-2.0
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+"""
+import os
+import pathlib
+from argparse import ArgumentDefaultsHelpFormatter, ArgumentParser
+
+import numpy as np
+import torch
+import torchvision.transforms as TF
+from PIL import Image
+from scipy import linalg
+from torch.nn.functional import adaptive_avg_pool2d
+
+try:
+ from tqdm import tqdm
+except ImportError:
+ # If tqdm is not available, provide a mock version of it
+ def tqdm(x):
+ return x
+
+class ImagePathDataset(torch.utils.data.Dataset):
+ def __init__(self, files, transforms=None):
+ self.files = files
+ self.transforms = transforms
+
+ def __len__(self):
+ return len(self.files)
+
+ def __getitem__(self, i):
+ path = self.files[i]
+ img = Image.open(path).convert('RGB')
+ if self.transforms is not None:
+ img = self.transforms(img)
+ return img
+
+
+def get_activations(files, model, batch_size=50, dims=2048, device='cpu',
+ num_workers=1):
+ """Calculates the activations of the pool_3 layer for all images.
+ Params:
+ -- files : List of image files paths
+ -- model : Instance of inception model
+ -- batch_size : Batch size of images for the model to process at once.
+ Make sure that the number of samples is a multiple of
+ the batch size, otherwise some samples are ignored. This
+ behavior is retained to match the original FID score
+ implementation.
+ -- dims : Dimensionality of features returned by Inception
+ -- device : Device to run calculations
+ -- num_workers : Number of parallel dataloader workers
+ Returns:
+ -- A numpy array of dimension (num images, dims) that contains the
+ activations of the given tensor when feeding inception with the
+ query tensor.
+ """
+ model.eval()
+
+ if batch_size > len(files):
+ print(('Warning: batch size is bigger than the data size. '
+ 'Setting batch size to data size'))
+ batch_size = len(files)
+
+ dataset = ImagePathDataset(files, transforms=TF.ToTensor())
+ dataloader = torch.utils.data.DataLoader(dataset,
+ batch_size=batch_size,
+ shuffle=False,
+ drop_last=False,
+ num_workers=num_workers)
+
+ pred_arr = np.empty((len(files), dims))
+
+ start_idx = 0
+
+ for batch in tqdm(dataloader):
+ batch = batch.to(device)
+
+ with torch.no_grad():
+ pred = model(batch)[0]
+
+ # If model output is not scalar, apply global spatial average pooling.
+ # This happens if you choose a dimensionality not equal 2048.
+ if pred.size(2) != 1 or pred.size(3) != 1:
+ pred = adaptive_avg_pool2d(pred, output_size=(1, 1))
+
+ pred = pred.squeeze(3).squeeze(2).cpu().numpy()
+
+ pred_arr[start_idx:start_idx + pred.shape[0]] = pred
+
+ start_idx = start_idx + pred.shape[0]
+
+ return pred_arr
+
+
+def calculate_frechet_distance(mu1, sigma1, mu2, sigma2, eps=1e-6):
+ """Numpy implementation of the Frechet Distance.
+ The Frechet distance between two multivariate Gaussians X_1 ~ N(mu_1, C_1)
+ and X_2 ~ N(mu_2, C_2) is
+ d^2 = ||mu_1 - mu_2||^2 + Tr(C_1 + C_2 - 2*sqrt(C_1*C_2)).
+ Stable version by Dougal J. Sutherland.
+ Params:
+ -- mu1 : Numpy array containing the activations of a layer of the
+ inception net (like returned by the function 'get_predictions')
+ for generated samples.
+ -- mu2 : The sample mean over activations, precalculated on an
+ representative data set.
+ -- sigma1: The covariance matrix over activations for generated samples.
+ -- sigma2: The covariance matrix over activations, precalculated on an
+ representative data set.
+ Returns:
+ -- : The Frechet Distance.
+ """
+
+ mu1 = np.atleast_1d(mu1)
+ mu2 = np.atleast_1d(mu2)
+
+ sigma1 = np.atleast_2d(sigma1)
+ sigma2 = np.atleast_2d(sigma2)
+
+ assert mu1.shape == mu2.shape, \
+ 'Training and test mean vectors have different lengths'
+ assert sigma1.shape == sigma2.shape, \
+ 'Training and test covariances have different dimensions'
+
+ diff = mu1 - mu2
+
+ # Product might be almost singular
+ covmean, _ = linalg.sqrtm(sigma1.dot(sigma2), disp=False)
+ if not np.isfinite(covmean).all():
+ msg = ('fid calculation produces singular product; '
+ 'adding %s to diagonal of cov estimates') % eps
+ print(msg)
+ offset = np.eye(sigma1.shape[0]) * eps
+ covmean = linalg.sqrtm((sigma1 + offset).dot(sigma2 + offset))
+
+ # Numerical error might give slight imaginary component
+ if np.iscomplexobj(covmean):
+ if not np.allclose(np.diagonal(covmean).imag, 0, atol=1e-3):
+ m = np.max(np.abs(covmean.imag))
+ raise ValueError('Imaginary component {}'.format(m))
+ covmean = covmean.real
+
+ tr_covmean = np.trace(covmean)
+
+ return (diff.dot(diff) + np.trace(sigma1)
+ + np.trace(sigma2) - 2 * tr_covmean)
+
+
+def calculate_activation_statistics(files, model, batch_size=50, dims=2048,
+ device='cpu', num_workers=1):
+ """Calculation of the statistics used by the FID.
+ Params:
+ -- files : List of image files paths
+ -- model : Instance of inception model
+ -- batch_size : The images numpy array is split into batches with
+ batch size batch_size. A reasonable batch size
+ depends on the hardware.
+ -- dims : Dimensionality of features returned by Inception
+ -- device : Device to run calculations
+ -- num_workers : Number of parallel dataloader workers
+ Returns:
+ -- mu : The mean over samples of the activations of the pool_3 layer of
+ the inception model.
+ -- sigma : The covariance matrix of the activations of the pool_3 layer of
+ the inception model.
+ """
+ act = get_activations(files, model, batch_size, dims, device, num_workers)
+ mu = np.mean(act, axis=0)
+ sigma = np.cov(act, rowvar=False)
+ return mu, sigma
diff --git a/lidm/eval/metric_utils.py b/lidm/eval/metric_utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..30210fc653842b19956a038385a007dccf7993e0
--- /dev/null
+++ b/lidm/eval/metric_utils.py
@@ -0,0 +1,458 @@
+"""
+@Author: Haoxi Ran
+@Date: 01/03/2024
+@Citation: Towards Realistic Scene Generation with LiDAR Diffusion Models
+
+"""
+
+import math
+from itertools import repeat
+from typing import List, Tuple, Union
+import numpy as np
+import torch
+
+from . import build_model, VOXEL_SIZE, MODALITY2MODEL, MODAL2BATCHSIZE, DATASET_CONFIG, AGG_TYPE, NUM_SECTORS, \
+ TYPE2DATASET, DATA_CONFIG
+
+try:
+ from torchsparse import SparseTensor, PointTensor
+ from torchsparse.utils.collate import sparse_collate_fn
+ from .modules.chamfer3D.dist_chamfer_3D import chamfer_3DDist
+ from .modules.chamfer2D.dist_chamfer_2D import chamfer_2DDist
+ from .modules.emd.emd_module import emdModule
+except:
+ print(
+ 'To install torchsparse 1.4.0, please refer to https://github.com/mit-han-lab/torchsparse/tree/74099d10a51c71c14318bce63d6421f698b24f24')
+
+
+def ravel_hash(x: np.ndarray) -> np.ndarray:
+ assert x.ndim == 2, x.shape
+
+ x = x - np.min(x, axis=0)
+ x = x.astype(np.uint64, copy=False)
+ xmax = np.max(x, axis=0).astype(np.uint64) + 1
+
+ h = np.zeros(x.shape[0], dtype=np.uint64)
+ for k in range(x.shape[1] - 1):
+ h += x[:, k]
+ h *= xmax[k + 1]
+ h += x[:, -1]
+ return h
+
+
+def sparse_quantize(coords, voxel_size: Union[float, Tuple[float, ...]] = 1, *, return_index: bool = False,
+ return_inverse: bool = False) -> List[np.ndarray]:
+ """
+ Modified based on https://github.com/mit-han-lab/torchsparse/blob/462dea4a701f87a7545afb3616bf2cf53dd404f3/torchsparse/utils/quantize.py
+
+ """
+ if isinstance(voxel_size, (float, int)):
+ voxel_size = tuple(repeat(voxel_size, coords.shape[1]))
+ assert isinstance(voxel_size, tuple) and len(voxel_size) in [2, 3] # support 2D and 3D coordinates only
+
+ voxel_size = np.array(voxel_size)
+ coords = np.floor(coords / voxel_size).astype(np.int32)
+
+ _, indices, inverse_indices = np.unique(
+ ravel_hash(coords), return_index=True, return_inverse=True
+ )
+ coords = coords[indices]
+
+ outputs = [coords]
+ if return_index:
+ outputs += [indices]
+ if return_inverse:
+ outputs += [inverse_indices]
+ return outputs[0] if len(outputs) == 1 else outputs
+
+
+def pcd2range(pcd, size, fov, depth_range, remission=None, labels=None, **kwargs):
+ # laser parameters
+ fov_up = fov[0] / 180.0 * np.pi # field of view up in rad
+ fov_down = fov[1] / 180.0 * np.pi # field of view down in rad
+ fov_range = abs(fov_down) + abs(fov_up) # get field of view total in rad
+
+ # get depth (distance) of all points
+ depth = np.linalg.norm(pcd, 2, axis=1)
+
+ # mask points out of range
+ mask = np.logical_and(depth > depth_range[0], depth < depth_range[1])
+ depth, pcd = depth[mask], pcd[mask]
+
+ # get scan components
+ scan_x, scan_y, scan_z = pcd[:, 0], pcd[:, 1], pcd[:, 2]
+
+ # get angles of all points
+ yaw = -np.arctan2(scan_y, scan_x)
+ pitch = np.arcsin(scan_z / depth)
+
+ # get projections in image coords
+ proj_x = 0.5 * (yaw / np.pi + 1.0) # in [0.0, 1.0]
+ proj_y = 1.0 - (pitch + abs(fov_down)) / fov_range # in [0.0, 1.0]
+
+ # scale to image size using angular resolution
+ proj_x *= size[1] # in [0.0, W]
+ proj_y *= size[0] # in [0.0, H]
+
+ # round and clamp for use as index
+ proj_x = np.maximum(0, np.minimum(size[1] - 1, np.floor(proj_x))).astype(np.int32) # in [0,W-1]
+ proj_y = np.maximum(0, np.minimum(size[0] - 1, np.floor(proj_y))).astype(np.int32) # in [0,H-1]
+
+ # order in decreasing depth
+ order = np.argsort(depth)[::-1]
+ proj_x, proj_y = proj_x[order], proj_y[order]
+
+ # project depth
+ depth = depth[order]
+ proj_range = np.full(size, -1, dtype=np.float32)
+ proj_range[proj_y, proj_x] = depth
+
+ # project point feature
+ if remission is not None:
+ remission = remission[mask][order]
+ proj_feature = np.full(size, -1, dtype=np.float32)
+ proj_feature[proj_y, proj_x] = remission
+ elif labels is not None:
+ labels = labels[mask][order]
+ proj_feature = np.full(size, 0, dtype=np.float32)
+ proj_feature[proj_y, proj_x] = labels
+ else:
+ proj_feature = None
+
+ return proj_range, proj_feature
+
+
+def range2xyz(range_img, fov, depth_range, depth_scale, log_scale=True, **kwargs):
+ # laser parameters
+ size = range_img.shape
+ fov_up = fov[0] / 180.0 * np.pi # field of view up in rad
+ fov_down = fov[1] / 180.0 * np.pi # field of view down in rad
+ fov_range = abs(fov_down) + abs(fov_up) # get field of view total in rad
+
+ # inverse transform from depth
+ if log_scale:
+ depth = (np.exp2(range_img * depth_scale) - 1)
+ else:
+ depth = range_img
+
+ scan_x, scan_y = np.meshgrid(np.arange(size[1]), np.arange(size[0]))
+ scan_x = scan_x.astype(np.float64) / size[1]
+ scan_y = scan_y.astype(np.float64) / size[0]
+
+ yaw = np.pi * (scan_x * 2 - 1)
+ pitch = (1.0 - scan_y) * fov_range - abs(fov_down)
+
+ xyz = -np.ones((3, *size))
+ xyz[0] = np.cos(yaw) * np.cos(pitch) * depth
+ xyz[1] = -np.sin(yaw) * np.cos(pitch) * depth
+ xyz[2] = np.sin(pitch) * depth
+
+ # mask out invalid points
+ mask = np.logical_and(depth > depth_range[0], depth < depth_range[1])
+ xyz[:, ~mask] = -1
+
+ return xyz
+
+
+def pcd2voxel(pcd):
+ pcd_voxel = np.round(pcd / VOXEL_SIZE)
+ pcd_voxel = pcd_voxel - pcd_voxel.min(0, keepdims=1)
+ feat = np.concatenate((pcd, -np.ones((pcd.shape[0], 1))), axis=1) # -1 for remission placeholder
+ _, inds, inverse_map = sparse_quantize(pcd_voxel, 1, return_index=True, return_inverse=True)
+
+ feat = torch.FloatTensor(feat[inds])
+ pcd_voxel = torch.LongTensor(pcd_voxel[inds])
+ lidar = SparseTensor(feat, pcd_voxel)
+ output = {'lidar': lidar}
+ return output
+
+
+def pcd2voxel_full(data_type, *args):
+ config = DATA_CONFIG[data_type]
+ x_range, y_range, z_range = config['x'], config['y'], config['z']
+ vol_shape = (math.ceil((x_range[1] - x_range[0]) / VOXEL_SIZE), math.ceil((y_range[1] - y_range[0]) / VOXEL_SIZE),
+ math.ceil((z_range[1] - z_range[0]) / VOXEL_SIZE))
+ min_bound = (math.ceil((x_range[0]) / VOXEL_SIZE), math.ceil((y_range[0]) / VOXEL_SIZE),
+ math.ceil((z_range[0]) / VOXEL_SIZE))
+
+ output = tuple()
+ for data in args:
+ volume_list = []
+ for pcd in data:
+ # mask out invalid points
+ mask_x = np.logical_and(pcd[:, 0] > x_range[0], pcd[:, 0] < x_range[1])
+ mask_y = np.logical_and(pcd[:, 1] > y_range[0], pcd[:, 1] < y_range[1])
+ mask_z = np.logical_and(pcd[:, 2] > z_range[0], pcd[:, 2] < z_range[1])
+ mask = mask_x & mask_y & mask_z
+ pcd = pcd[mask]
+
+ # voxelize
+ pcd_voxel = np.floor(pcd / VOXEL_SIZE)
+ _, indices, inverse_map = sparse_quantize(pcd_voxel, 1, return_index=True, return_inverse=True)
+ pcd_voxel = pcd_voxel[indices]
+ pcd_voxel = (pcd_voxel - min_bound).astype(np.int32)
+
+ # 2D bev grid
+ vol = np.zeros(vol_shape, dtype=np.float32)
+ vol[pcd_voxel[:, 0], pcd_voxel[:, 1], pcd_voxel[:, 2]] = 1
+ volume_list.append(vol)
+ output += (volume_list,)
+ return output
+
+
+# def pcd2bev_full(data_type, *args, voxel_size=VOXEL_SIZE):
+# config = DATA_CONFIG[data_type]
+# x_range, y_range = config['x'], config['y']
+# vol_shape = (math.ceil((x_range[1] - x_range[0]) / voxel_size), math.ceil((y_range[1] - y_range[0]) / voxel_size))
+# min_bound = (math.ceil((x_range[0]) / voxel_size), math.ceil((y_range[0]) / voxel_size))
+#
+# output = tuple()
+# for data in args:
+# volume_list = []
+# for pcd in data:
+# # mask out invalid points
+# mask_x = np.logical_and(pcd[:, 0] > x_range[0], pcd[:, 0] < x_range[1])
+# mask_y = np.logical_and(pcd[:, 1] > y_range[0], pcd[:, 1] < y_range[1])
+# mask = mask_x & mask_y
+# pcd = pcd[mask][:, :2] # keep x,y coord
+#
+# # voxelize
+# pcd_voxel = np.floor(pcd / voxel_size)
+# _, indices, inverse_map = sparse_quantize(pcd_voxel, 1, return_index=True, return_inverse=True)
+# pcd_voxel = pcd_voxel[indices]
+# pcd_voxel = (pcd_voxel - min_bound).astype(np.int32)
+#
+# # 2D bev grid
+# vol = np.zeros(vol_shape, dtype=np.float32)
+# vol[pcd_voxel[:, 0], pcd_voxel[:, 1]] = 1
+# volume_list.append(vol)
+# output += (volume_list,)
+# return output
+
+
+def pcd2bev_sum(data_type, *args, voxel_size=VOXEL_SIZE):
+ config = DATA_CONFIG[data_type]
+ x_range, y_range = config['x'], config['y']
+ vol_shape = (math.ceil((x_range[1] - x_range[0]) / voxel_size), math.ceil((y_range[1] - y_range[0]) / voxel_size))
+ min_bound = (math.ceil((x_range[0]) / voxel_size), math.ceil((y_range[0]) / voxel_size))
+
+ output = tuple()
+ for data in args:
+ volume_sum = np.zeros(vol_shape, np.float32)
+ for pcd in data:
+ # mask out invalid points
+ mask_x = np.logical_and(pcd[:, 0] > x_range[0], pcd[:, 0] < x_range[1])
+ mask_y = np.logical_and(pcd[:, 1] > y_range[0], pcd[:, 1] < y_range[1])
+ mask = mask_x & mask_y
+ pcd = pcd[mask][:, :2] # keep x,y coord
+
+ # voxelize
+ pcd_voxel = np.floor(pcd / voxel_size)
+ _, indices, inverse_map = sparse_quantize(pcd_voxel, 1, return_index=True, return_inverse=True)
+ pcd_voxel = pcd_voxel[indices]
+ pcd_voxel = (pcd_voxel - min_bound).astype(np.int32)
+
+ # summation
+ volume_sum[pcd_voxel[:, 0], pcd_voxel[:, 1]] += 1.
+ output += (volume_sum,)
+ return output
+
+
+def pcd2bev_bin(data_type, *args, voxel_size=0.5):
+ config = DATA_CONFIG[data_type]
+ x_range, y_range = config['x'], config['y']
+ vol_shape = (math.ceil((x_range[1] - x_range[0]) / voxel_size), math.ceil((y_range[1] - y_range[0]) / voxel_size))
+ min_bound = (math.ceil((x_range[0]) / voxel_size), math.ceil((y_range[0]) / voxel_size))
+
+ output = tuple()
+ for data in args:
+ pcd_list = []
+ for pcd in data:
+ # mask out invalid points
+ mask_x = np.logical_and(pcd[:, 0] > x_range[0], pcd[:, 0] < x_range[1])
+ mask_y = np.logical_and(pcd[:, 1] > y_range[0], pcd[:, 1] < y_range[1])
+ mask = mask_x & mask_y
+ pcd = pcd[mask][:, :2] # keep x,y coord
+
+ # voxelize
+ pcd_voxel = np.floor(pcd / voxel_size)
+ _, indices, inverse_map = sparse_quantize(pcd_voxel, 1, return_index=True, return_inverse=True)
+ pcd_voxel = pcd_voxel[indices]
+ pcd_voxel = ((pcd_voxel - min_bound) / vol_shape).astype(np.float32)
+ pcd_list.append(pcd_voxel)
+ output += (pcd_list,)
+ return output
+
+
+def bev_sample(data_type, *args, voxel_size=0.5):
+ config = DATA_CONFIG[data_type]
+ x_range, y_range = config['x'], config['y']
+
+ output = tuple()
+ for data in args:
+ pcd_list = []
+ for pcd in data:
+ # mask out invalid points
+ mask_x = np.logical_and(pcd[:, 0] > x_range[0], pcd[:, 0] < x_range[1])
+ mask_y = np.logical_and(pcd[:, 1] > y_range[0], pcd[:, 1] < y_range[1])
+ mask = mask_x & mask_y
+ pcd = pcd[mask][:, :2] # keep x,y coord
+
+ # voxelize
+ pcd_voxel = np.floor(pcd / voxel_size)
+ _, indices, inverse_map = sparse_quantize(pcd_voxel, 1, return_index=True, return_inverse=True)
+ pcd = pcd[indices]
+ pcd_list.append(pcd)
+ output += (pcd_list,)
+ return output
+
+
+def preprocess_pcd(pcd, **kwargs):
+ depth = np.linalg.norm(pcd, 2, axis=1)
+ mask = np.logical_and(depth > kwargs['depth_range'][0], depth < kwargs['depth_range'][1])
+ pcd = pcd[mask]
+ return pcd
+
+
+def preprocess_range(pcd, **kwargs):
+ depth_img = pcd2range(pcd, **kwargs)[0]
+ xyz_img = range2xyz(depth_img, log_scale=False, **kwargs)
+ depth_img = depth_img[None]
+ img = np.vstack([depth_img, xyz_img])
+ return img
+
+
+def batch2list(batch_dict, agg_type='depth', **kwargs):
+ """
+ Aggregation Type: Default 'depth', ['all', 'sector', 'depth']
+ """
+ output_list = []
+ batch_indices = batch_dict['batch_indices']
+ for b_idx in range(batch_indices.max() + 1):
+ # avg all
+ if agg_type == 'all':
+ logits = batch_dict['logits'][batch_indices == b_idx].mean(0)
+
+ # avg on sectors
+ elif agg_type == 'sector':
+ logits = batch_dict['logits'][batch_indices == b_idx]
+ coords = batch_dict['coords'][batch_indices == b_idx].float()
+ coords = coords - coords.mean(0)
+ angle = torch.atan2(coords[:, 1], coords[:, 0]) # [-pi, pi]
+ sector_range = torch.linspace(-np.pi - 1e-4, np.pi + 1e-4, NUM_SECTORS + 1)
+ logits_list = []
+ for i in range(NUM_SECTORS):
+ sector_indices = torch.where((angle >= sector_range[i]) & (angle < sector_range[i + 1]))[0]
+ sector_logits = logits[sector_indices].mean(0)
+ sector_logits = torch.nan_to_num(sector_logits, 0.)
+ logits_list.append(sector_logits)
+ logits = torch.cat(logits_list) # dim: 768
+
+ # avg by depth
+ elif agg_type == 'depth':
+ logits = batch_dict['logits'][batch_indices == b_idx]
+ coords = batch_dict['coords'][batch_indices == b_idx].float()
+ coords = coords - coords.mean(0)
+ bev_depth = torch.norm(coords, dim=-1) * VOXEL_SIZE
+ sector_range = torch.linspace(kwargs['depth_range'][0] + 3, kwargs['depth_range'][1], NUM_SECTORS + 1)
+ sector_range[0] = 0.
+ logits_list = []
+ for i in range(NUM_SECTORS):
+ sector_indices = torch.where((bev_depth >= sector_range[i]) & (bev_depth < sector_range[i + 1]))[0]
+ sector_logits = logits[sector_indices].mean(0)
+ sector_logits = torch.nan_to_num(sector_logits, 0.)
+ logits_list.append(sector_logits)
+ logits = torch.cat(logits_list) # dim: 768
+
+ else:
+ raise NotImplementedError
+
+ output_list.append(logits.detach().cpu().numpy())
+ return output_list
+
+
+def compute_logits(data_type, modality, *args):
+ assert data_type in ['32', '64']
+ assert modality in ['range', 'voxel', 'point_voxel']
+ is_voxel = 'voxel' in modality
+ dataset_name = TYPE2DATASET[data_type]
+ dataset_config = DATASET_CONFIG[dataset_name]
+ bs = MODAL2BATCHSIZE[modality]
+
+ model = build_model(dataset_name, MODALITY2MODEL[modality], device='cuda')
+
+ output = tuple()
+ for data in args:
+ all_logits_list = []
+ for i in range(math.ceil(len(data) / bs)):
+ batch = data[i * bs:(i + 1) * bs]
+ if is_voxel:
+ batch = [pcd2voxel(preprocess_pcd(pcd, **dataset_config)) for pcd in batch]
+ batch = sparse_collate_fn(batch)
+ batch = {k: v.cuda() if isinstance(v, (torch.Tensor, SparseTensor, PointTensor)) else v for k, v in
+ batch.items()}
+ with torch.no_grad():
+ batch_out = model(batch, return_final_logits=True)
+ batch_out = batch2list(batch_out, AGG_TYPE, **dataset_config)
+ all_logits_list.extend(batch_out)
+ else:
+ batch = [preprocess_range(pcd, **dataset_config) for pcd in batch]
+ batch = torch.from_numpy(np.stack(batch)).float().cuda()
+ with torch.no_grad():
+ batch_out = model(batch, return_final_logits=True, agg_type=AGG_TYPE)
+ all_logits_list.append(batch_out)
+ if is_voxel:
+ all_logits = np.stack(all_logits_list)
+ else:
+ all_logits = np.vstack(all_logits_list)
+ output += (all_logits,)
+
+ del model, batch, batch_out
+ torch.cuda.empty_cache()
+ return output
+
+
+def compute_pairwise_cd(x, y, module=None):
+ if module is None:
+ module = chamfer_3DDist()
+ if x.ndim == 2 and y.ndim == 2:
+ x, y = x[None], y[None]
+ x, y = torch.from_numpy(x).cuda(), torch.from_numpy(y).cuda()
+ dist1, dist2, _, _ = module(x, y)
+ dist = (dist1.mean() + dist2.mean()) / 2
+ return dist.item()
+
+
+def compute_pairwise_cd_batch(reference, samples):
+ ndim = reference.ndim
+ assert ndim in [2, 3]
+ module = chamfer_3DDist() if ndim == 3 else chamfer_2DDist()
+ len_r, len_s = reference.shape[0], [s.shape[0] for s in samples]
+ max_len = max([len_r] + len_s)
+ reference = torch.from_numpy(
+ np.vstack([reference, np.ones((max_len - reference.shape[0], ndim), dtype=np.float32) * 1e6])).cuda()
+ samples = [np.vstack([s, np.ones((max_len - s.shape[0], ndim), dtype=np.float32) * 1e6]) for s in samples]
+ samples = torch.from_numpy(np.stack(samples)).cuda()
+ reference = reference.expand_as(samples)
+ dist_r, dist_s, _, _ = module(reference, samples)
+
+ results = []
+ for i in range(samples.shape[0]):
+ dist1, dist2, len1, len2 = dist_r[i], dist_s[i], len_r, len_s[i]
+ dist = (dist1[:len1].mean() + dist2[:len2].mean()) / 2.
+ results.append(dist.item())
+ return results
+
+
+def compute_pairwise_emd(x, y, module=None):
+ if module is None:
+ module = emdModule()
+ n_points = min(x.shape[0], y.shape[0])
+ n_points = n_points - n_points % 1024
+ x, y = x[:n_points], y[:n_points]
+ if x.ndim == 2 and y.ndim == 2:
+ x, y = x[None], y[None]
+ x, y = torch.from_numpy(x).cuda(), torch.from_numpy(y).cuda()
+ dist, _ = module(x, y, 0.005, 50)
+ dist = torch.sqrt(dist).mean()
+ return dist.item()
diff --git a/lidm/eval/models/__init__.py b/lidm/eval/models/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/lidm/eval/models/minkowskinet/__init__.py b/lidm/eval/models/minkowskinet/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/lidm/eval/models/minkowskinet/model.py b/lidm/eval/models/minkowskinet/model.py
new file mode 100644
index 0000000000000000000000000000000000000000..daa36a5009e96f3331653341c6b185627d688a19
--- /dev/null
+++ b/lidm/eval/models/minkowskinet/model.py
@@ -0,0 +1,141 @@
+import torch
+import torch.nn as nn
+
+try:
+ import torchsparse
+ import torchsparse.nn as spnn
+ from ..ts import basic_blocks
+except ImportError:
+ raise Exception('Required ts lib. Reference: https://github.com/mit-han-lab/torchsparse/tree/v1.4.0')
+
+
+class Model(nn.Module):
+ def __init__(self, config):
+ super().__init__()
+
+ cr = config.model_params.cr
+ cs = config.model_params.layer_num
+ cs = [int(cr * x) for x in cs]
+
+ self.pres = self.vres = config.model_params.voxel_size
+ self.num_classes = config.model_params.num_class
+
+ self.stem = nn.Sequential(
+ spnn.Conv3d(config.model_params.input_dims, cs[0], kernel_size=3, stride=1),
+ spnn.BatchNorm(cs[0]), spnn.ReLU(True),
+ spnn.Conv3d(cs[0], cs[0], kernel_size=3, stride=1),
+ spnn.BatchNorm(cs[0]), spnn.ReLU(True))
+
+ self.stage1 = nn.Sequential(
+ basic_blocks.BasicConvolutionBlock(cs[0], cs[0], ks=2, stride=2, dilation=1),
+ basic_blocks.ResidualBlock(cs[0], cs[1], ks=3, stride=1, dilation=1),
+ basic_blocks.ResidualBlock(cs[1], cs[1], ks=3, stride=1, dilation=1),
+ )
+
+ self.stage2 = nn.Sequential(
+ basic_blocks.BasicConvolutionBlock(cs[1], cs[1], ks=2, stride=2, dilation=1),
+ basic_blocks.ResidualBlock(cs[1], cs[2], ks=3, stride=1, dilation=1),
+ basic_blocks.ResidualBlock(cs[2], cs[2], ks=3, stride=1, dilation=1),
+ )
+
+ self.stage3 = nn.Sequential(
+ basic_blocks.BasicConvolutionBlock(cs[2], cs[2], ks=2, stride=2, dilation=1),
+ basic_blocks.ResidualBlock(cs[2], cs[3], ks=3, stride=1, dilation=1),
+ basic_blocks.ResidualBlock(cs[3], cs[3], ks=3, stride=1, dilation=1),
+ )
+
+ self.stage4 = nn.Sequential(
+ basic_blocks.BasicConvolutionBlock(cs[3], cs[3], ks=2, stride=2, dilation=1),
+ basic_blocks.ResidualBlock(cs[3], cs[4], ks=3, stride=1, dilation=1),
+ basic_blocks.ResidualBlock(cs[4], cs[4], ks=3, stride=1, dilation=1),
+ )
+
+ self.up1 = nn.ModuleList([
+ basic_blocks.BasicDeconvolutionBlock(cs[4], cs[5], ks=2, stride=2),
+ nn.Sequential(
+ basic_blocks.ResidualBlock(cs[5] + cs[3], cs[5], ks=3, stride=1,
+ dilation=1),
+ basic_blocks.ResidualBlock(cs[5], cs[5], ks=3, stride=1, dilation=1),
+ )
+ ])
+
+ self.up2 = nn.ModuleList([
+ basic_blocks.BasicDeconvolutionBlock(cs[5], cs[6], ks=2, stride=2),
+ nn.Sequential(
+ basic_blocks.ResidualBlock(cs[6] + cs[2], cs[6], ks=3, stride=1,
+ dilation=1),
+ basic_blocks.ResidualBlock(cs[6], cs[6], ks=3, stride=1, dilation=1),
+ )
+ ])
+
+ self.up3 = nn.ModuleList([
+ basic_blocks.BasicDeconvolutionBlock(cs[6], cs[7], ks=2, stride=2),
+ nn.Sequential(
+ basic_blocks.ResidualBlock(cs[7] + cs[1], cs[7], ks=3, stride=1,
+ dilation=1),
+ basic_blocks.ResidualBlock(cs[7], cs[7], ks=3, stride=1, dilation=1),
+ )
+ ])
+
+ self.up4 = nn.ModuleList([
+ basic_blocks.BasicDeconvolutionBlock(cs[7], cs[8], ks=2, stride=2),
+ nn.Sequential(
+ basic_blocks.ResidualBlock(cs[8] + cs[0], cs[8], ks=3, stride=1,
+ dilation=1),
+ basic_blocks.ResidualBlock(cs[8], cs[8], ks=3, stride=1, dilation=1),
+ )
+ ])
+
+ self.classifier = nn.Sequential(nn.Linear(cs[8], self.num_classes))
+
+ self.weight_initialization()
+ self.dropout = nn.Dropout(0.3, True)
+
+ def weight_initialization(self):
+ for m in self.modules():
+ if isinstance(m, nn.BatchNorm1d):
+ nn.init.constant_(m.weight, 1)
+ nn.init.constant_(m.bias, 0)
+
+ def forward(self, data_dict, return_logits=False, return_final_logits=False):
+ x = data_dict['lidar']
+ x.C = x.C.int()
+
+ x0 = self.stem(x)
+ x1 = self.stage1(x0)
+ x2 = self.stage2(x1)
+ x3 = self.stage3(x2)
+ x4 = self.stage4(x3)
+
+ if return_logits:
+ output_dict = dict()
+ output_dict['logits'] = x4.F
+ output_dict['batch_indices'] = x4.C[:, -1]
+ return output_dict
+
+ y1 = self.up1[0](x4)
+ y1 = torchsparse.cat([y1, x3])
+ y1 = self.up1[1](y1)
+
+ y2 = self.up2[0](y1)
+ y2 = torchsparse.cat([y2, x2])
+ y2 = self.up2[1](y2)
+
+ y3 = self.up3[0](y2)
+ y3 = torchsparse.cat([y3, x1])
+ y3 = self.up3[1](y3)
+
+ y4 = self.up4[0](y3)
+ y4 = torchsparse.cat([y4, x0])
+ y4 = self.up4[1](y4)
+ if return_final_logits:
+ output_dict = dict()
+ output_dict['logits'] = y4.F
+ output_dict['coords'] = y4.C[:, :3]
+ output_dict['batch_indices'] = y4.C[:, -1]
+ return output_dict
+
+ output = self.classifier(y4.F)
+ data_dict['output'] = output.F
+
+ return data_dict
diff --git a/lidm/eval/models/rangenet/__init__.py b/lidm/eval/models/rangenet/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/lidm/eval/models/rangenet/model.py b/lidm/eval/models/rangenet/model.py
new file mode 100644
index 0000000000000000000000000000000000000000..752fae9effd476a6bff255e0063675d1cc2f72e2
--- /dev/null
+++ b/lidm/eval/models/rangenet/model.py
@@ -0,0 +1,372 @@
+#!/usr/bin/env python3
+# This file is covered by the LICENSE file in the root of this project.
+from collections import OrderedDict
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class BasicBlock(nn.Module):
+ def __init__(self, inplanes, planes, bn_d=0.1):
+ super(BasicBlock, self).__init__()
+ self.conv1 = nn.Conv2d(inplanes, planes[0], kernel_size=1,
+ stride=1, padding=0, bias=False)
+ self.bn1 = nn.BatchNorm2d(planes[0], momentum=bn_d)
+ self.relu1 = nn.LeakyReLU(0.1)
+ self.conv2 = nn.Conv2d(planes[0], planes[1], kernel_size=3,
+ stride=1, padding=1, bias=False)
+ self.bn2 = nn.BatchNorm2d(planes[1], momentum=bn_d)
+ self.relu2 = nn.LeakyReLU(0.1)
+
+ def forward(self, x):
+ residual = x
+
+ out = self.conv1(x)
+ out = self.bn1(out)
+ out = self.relu1(out)
+
+ out = self.conv2(out)
+ out = self.bn2(out)
+ out = self.relu2(out)
+
+ out += residual
+ return out
+
+
+# ******************************************************************************
+
+# number of layers per model
+model_blocks = {
+ 21: [1, 1, 2, 2, 1],
+ 53: [1, 2, 8, 8, 4],
+}
+
+
+class Backbone(nn.Module):
+ """
+ Class for DarknetSeg. Subclasses PyTorch's own "nn" module
+ """
+
+ def __init__(self, params):
+ super(Backbone, self).__init__()
+ self.use_range = params["input_depth"]["range"]
+ self.use_xyz = params["input_depth"]["xyz"]
+ self.use_remission = params["input_depth"]["remission"]
+ self.drop_prob = params["dropout"]
+ self.bn_d = params["bn_d"]
+ self.OS = params["OS"]
+ self.layers = params["extra"]["layers"]
+
+ # input depth calc
+ self.input_depth = 0
+ self.input_idxs = []
+ if self.use_range:
+ self.input_depth += 1
+ self.input_idxs.append(0)
+ if self.use_xyz:
+ self.input_depth += 3
+ self.input_idxs.extend([1, 2, 3])
+ if self.use_remission:
+ self.input_depth += 1
+ self.input_idxs.append(4)
+
+ # stride play
+ self.strides = [2, 2, 2, 2, 2]
+ # check current stride
+ current_os = 1
+ for s in self.strides:
+ current_os *= s
+
+ # make the new stride
+ if self.OS > current_os:
+ print("Can't do OS, ", self.OS,
+ " because it is bigger than original ", current_os)
+ else:
+ # redo strides according to needed stride
+ for i, stride in enumerate(reversed(self.strides), 0):
+ if int(current_os) != self.OS:
+ if stride == 2:
+ current_os /= 2
+ self.strides[-1 - i] = 1
+ if int(current_os) == self.OS:
+ break
+
+ # check that darknet exists
+ assert self.layers in model_blocks.keys()
+
+ # generate layers depending on darknet type
+ self.blocks = model_blocks[self.layers]
+
+ # input layer
+ self.conv1 = nn.Conv2d(self.input_depth, 32, kernel_size=3,
+ stride=1, padding=1, bias=False)
+ self.bn1 = nn.BatchNorm2d(32, momentum=self.bn_d)
+ self.relu1 = nn.LeakyReLU(0.1)
+
+ # encoder
+ self.enc1 = self._make_enc_layer(BasicBlock, [32, 64], self.blocks[0],
+ stride=self.strides[0], bn_d=self.bn_d)
+ self.enc2 = self._make_enc_layer(BasicBlock, [64, 128], self.blocks[1],
+ stride=self.strides[1], bn_d=self.bn_d)
+ self.enc3 = self._make_enc_layer(BasicBlock, [128, 256], self.blocks[2],
+ stride=self.strides[2], bn_d=self.bn_d)
+ self.enc4 = self._make_enc_layer(BasicBlock, [256, 512], self.blocks[3],
+ stride=self.strides[3], bn_d=self.bn_d)
+ self.enc5 = self._make_enc_layer(BasicBlock, [512, 1024], self.blocks[4],
+ stride=self.strides[4], bn_d=self.bn_d)
+
+ # for a bit of fun
+ self.dropout = nn.Dropout2d(self.drop_prob)
+
+ # last channels
+ self.last_channels = 1024
+
+ # make layer useful function
+ def _make_enc_layer(self, block, planes, blocks, stride, bn_d=0.1):
+ layers = []
+
+ # downsample
+ layers.append(("conv", nn.Conv2d(planes[0], planes[1],
+ kernel_size=3,
+ stride=[1, stride], dilation=1,
+ padding=1, bias=False)))
+ layers.append(("bn", nn.BatchNorm2d(planes[1], momentum=bn_d)))
+ layers.append(("relu", nn.LeakyReLU(0.1)))
+
+ # blocks
+ inplanes = planes[1]
+ for i in range(0, blocks):
+ layers.append(("residual_{}".format(i),
+ block(inplanes, planes, bn_d)))
+
+ return nn.Sequential(OrderedDict(layers))
+
+ def run_layer(self, x, layer, skips, os):
+ y = layer(x)
+ if y.shape[2] < x.shape[2] or y.shape[3] < x.shape[3]:
+ skips[os] = x.detach()
+ os *= 2
+ x = y
+ return x, skips, os
+
+ def forward(self, x, return_logits=False, return_list=None):
+ # filter input
+ x = x[:, self.input_idxs]
+
+ # run cnn
+ # store for skip connections
+ skips = {}
+ out_dict = {}
+ os = 1
+
+ # first layer
+ x, skips, os = self.run_layer(x, self.conv1, skips, os)
+ x, skips, os = self.run_layer(x, self.bn1, skips, os)
+ x, skips, os = self.run_layer(x, self.relu1, skips, os)
+ if return_list and 'enc_0' in return_list:
+ out_dict['enc_0'] = x.detach().cpu() # 32, 64, 1024
+
+ # all encoder blocks with intermediate dropouts
+ x, skips, os = self.run_layer(x, self.enc1, skips, os)
+ if return_list and 'enc_1' in return_list:
+ out_dict['enc_1'] = x.detach().cpu() # 64, 64, 512
+ x, skips, os = self.run_layer(x, self.dropout, skips, os)
+
+ x, skips, os = self.run_layer(x, self.enc2, skips, os)
+ if return_list and 'enc_2' in return_list:
+ out_dict['enc_2'] = x.detach().cpu() # 128, 64, 256
+ x, skips, os = self.run_layer(x, self.dropout, skips, os)
+
+ x, skips, os = self.run_layer(x, self.enc3, skips, os)
+ if return_list and 'enc_3' in return_list:
+ out_dict['enc_3'] = x.detach().cpu() # 256, 64, 128
+ x, skips, os = self.run_layer(x, self.dropout, skips, os)
+
+ x, skips, os = self.run_layer(x, self.enc4, skips, os)
+ if return_list and 'enc_4' in return_list:
+ out_dict['enc_4'] = x.detach().cpu() # 512, 64, 64
+ x, skips, os = self.run_layer(x, self.dropout, skips, os)
+
+ x, skips, os = self.run_layer(x, self.enc5, skips, os)
+ if return_list and 'enc_5' in return_list:
+ out_dict['enc_5'] = x.detach().cpu() # 1024, 64, 32
+ if return_logits:
+ return x
+
+ x, skips, os = self.run_layer(x, self.dropout, skips, os)
+
+ if return_list is not None:
+ return x, skips, out_dict
+ return x, skips
+
+ def get_last_depth(self):
+ return self.last_channels
+
+ def get_input_depth(self):
+ return self.input_depth
+
+
+class Decoder(nn.Module):
+ """
+ Class for DarknetSeg. Subclasses PyTorch's own "nn" module
+ """
+
+ def __init__(self, params, OS=32, feature_depth=1024):
+ super(Decoder, self).__init__()
+ self.backbone_OS = OS
+ self.backbone_feature_depth = feature_depth
+ self.drop_prob = params["dropout"]
+ self.bn_d = params["bn_d"]
+ self.index = 0
+
+ # stride play
+ self.strides = [2, 2, 2, 2, 2]
+ # check current stride
+ current_os = 1
+ for s in self.strides:
+ current_os *= s
+ # redo strides according to needed stride
+ for i, stride in enumerate(self.strides):
+ if int(current_os) != self.backbone_OS:
+ if stride == 2:
+ current_os /= 2
+ self.strides[i] = 1
+ if int(current_os) == self.backbone_OS:
+ break
+
+ # decoder
+ self.dec5 = self._make_dec_layer(BasicBlock,
+ [self.backbone_feature_depth, 512],
+ bn_d=self.bn_d,
+ stride=self.strides[0])
+ self.dec4 = self._make_dec_layer(BasicBlock, [512, 256], bn_d=self.bn_d,
+ stride=self.strides[1])
+ self.dec3 = self._make_dec_layer(BasicBlock, [256, 128], bn_d=self.bn_d,
+ stride=self.strides[2])
+ self.dec2 = self._make_dec_layer(BasicBlock, [128, 64], bn_d=self.bn_d,
+ stride=self.strides[3])
+ self.dec1 = self._make_dec_layer(BasicBlock, [64, 32], bn_d=self.bn_d,
+ stride=self.strides[4])
+
+ # layer list to execute with skips
+ self.layers = [self.dec5, self.dec4, self.dec3, self.dec2, self.dec1]
+
+ # for a bit of fun
+ self.dropout = nn.Dropout2d(self.drop_prob)
+
+ # last channels
+ self.last_channels = 32
+
+ def _make_dec_layer(self, block, planes, bn_d=0.1, stride=2):
+ layers = []
+
+ # downsample
+ if stride == 2:
+ layers.append(("upconv", nn.ConvTranspose2d(planes[0], planes[1],
+ kernel_size=[1, 4], stride=[1, 2],
+ padding=[0, 1])))
+ else:
+ layers.append(("conv", nn.Conv2d(planes[0], planes[1],
+ kernel_size=3, padding=1)))
+ layers.append(("bn", nn.BatchNorm2d(planes[1], momentum=bn_d)))
+ layers.append(("relu", nn.LeakyReLU(0.1)))
+
+ # blocks
+ layers.append(("residual", block(planes[1], planes, bn_d)))
+
+ return nn.Sequential(OrderedDict(layers))
+
+ def run_layer(self, x, layer, skips, os):
+ feats = layer(x) # up
+ if feats.shape[-1] > x.shape[-1]:
+ os //= 2 # match skip
+ feats = feats + skips[os].detach() # add skip
+ x = feats
+ return x, skips, os
+
+ def forward(self, x, skips, return_logits=False, return_list=None):
+ os = self.backbone_OS
+ out_dict = {}
+
+ # run layers
+ x, skips, os = self.run_layer(x, self.dec5, skips, os)
+ if return_list and 'dec_4' in return_list:
+ out_dict['dec_4'] = x.detach().cpu() # 512, 64, 64
+ x, skips, os = self.run_layer(x, self.dec4, skips, os)
+ if return_list and 'dec_3' in return_list:
+ out_dict['dec_3'] = x.detach().cpu() # 256, 64, 128
+ x, skips, os = self.run_layer(x, self.dec3, skips, os)
+ if return_list and 'dec_2' in return_list:
+ out_dict['dec_2'] = x.detach().cpu() # 128, 64, 256
+ x, skips, os = self.run_layer(x, self.dec2, skips, os)
+ if return_list and 'dec_1' in return_list:
+ out_dict['dec_1'] = x.detach().cpu() # 64, 64, 512
+ x, skips, os = self.run_layer(x, self.dec1, skips, os)
+ if return_list and 'dec_0' in return_list:
+ out_dict['dec_0'] = x.detach().cpu() # 32, 64, 1024
+
+ logits = torch.clone(x).detach()
+ x = self.dropout(x)
+
+ if return_logits:
+ return x, logits
+ if return_list is not None:
+ return out_dict
+ return x
+
+ def get_last_depth(self):
+ return self.last_channels
+
+
+class Model(nn.Module):
+ def __init__(self, config):
+ super().__init__()
+ self.config = config
+ self.backbone = Backbone(params=self.config["backbone"])
+ self.decoder = Decoder(params=self.config["decoder"], OS=self.config["backbone"]["OS"],
+ feature_depth=self.backbone.get_last_depth())
+
+ def load_pretrained_weights(self, path):
+ w_dict = torch.load(path + "/backbone",
+ map_location=lambda storage, loc: storage)
+ self.backbone.load_state_dict(w_dict, strict=True)
+ w_dict = torch.load(path + "/segmentation_decoder",
+ map_location=lambda storage, loc: storage)
+ self.decoder.load_state_dict(w_dict, strict=True)
+
+ def forward(self, x, return_logits=False, return_final_logits=False, return_list=None, agg_type='depth'):
+ if return_logits:
+ logits = self.backbone(x, return_logits)
+ logits = F.adaptive_avg_pool2d(logits, (1, 1)).squeeze()
+ logits = torch.clone(logits).detach().cpu().numpy()
+ return logits
+ elif return_list is not None:
+ x, skips, enc_dict = self.backbone(x, return_list=return_list)
+ dec_dict = self.decoder(x, skips, return_list=return_list)
+ out_dict = {**enc_dict, **dec_dict}
+ return out_dict
+ elif return_final_logits:
+ assert agg_type in ['all', 'sector', 'depth']
+ y, skips = self.backbone(x)
+ y, logits = self.decoder(y, skips, True)
+
+ B, C, H, W = logits.shape
+ N = 16
+
+ # avg all
+ if agg_type == 'all':
+ logits = logits.mean([2, 3])
+ # avg in patch
+ elif agg_type == 'sector':
+ logits = logits.view(B, C, H, N, W // N).mean([2, 4]).reshape(B, -1)
+ # avg in row
+ elif agg_type == 'depth':
+ logits = logits.view(B, C, N, H // N, W).mean([3, 4]).reshape(B, -1)
+
+ logits = torch.clone(logits).detach().cpu().numpy()
+ return logits
+ else:
+ y, skips = self.backbone(x)
+ y = self.decoder(y, skips, False)
+ return y
diff --git a/lidm/eval/models/spvcnn/__init__.py b/lidm/eval/models/spvcnn/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/lidm/eval/models/spvcnn/model.py b/lidm/eval/models/spvcnn/model.py
new file mode 100644
index 0000000000000000000000000000000000000000..dd7793f8e81cef35f331c0c2c70062023999e83a
--- /dev/null
+++ b/lidm/eval/models/spvcnn/model.py
@@ -0,0 +1,179 @@
+import torch.nn as nn
+
+try:
+ import torchsparse
+ import torchsparse.nn as spnn
+ from torchsparse import PointTensor
+ from ..ts.utils import initial_voxelize, point_to_voxel, voxel_to_point
+ from ..ts import basic_blocks
+except ImportError:
+ raise Exception('Required torchsparse lib. Reference: https://github.com/mit-han-lab/torchsparse/tree/v1.4.0')
+
+
+class Model(nn.Module):
+ def __init__(self, config):
+ super().__init__()
+ cr = config.model_params.cr
+ cs = config.model_params.layer_num
+ cs = [int(cr * x) for x in cs]
+
+ self.pres = self.vres = config.model_params.voxel_size
+ self.num_classes = config.model_params.num_class
+
+ self.stem = nn.Sequential(
+ spnn.Conv3d(config.model_params.input_dims, cs[0], kernel_size=3, stride=1),
+ spnn.BatchNorm(cs[0]), spnn.ReLU(True),
+ spnn.Conv3d(cs[0], cs[0], kernel_size=3, stride=1),
+ spnn.BatchNorm(cs[0]), spnn.ReLU(True))
+
+ self.stage1 = nn.Sequential(
+ basic_blocks.BasicConvolutionBlock(cs[0], cs[0], ks=2, stride=2, dilation=1),
+ basic_blocks.ResidualBlock(cs[0], cs[1], ks=3, stride=1, dilation=1),
+ basic_blocks.ResidualBlock(cs[1], cs[1], ks=3, stride=1, dilation=1),
+ )
+
+ self.stage2 = nn.Sequential(
+ basic_blocks.BasicConvolutionBlock(cs[1], cs[1], ks=2, stride=2, dilation=1),
+ basic_blocks.ResidualBlock(cs[1], cs[2], ks=3, stride=1, dilation=1),
+ basic_blocks.ResidualBlock(cs[2], cs[2], ks=3, stride=1, dilation=1),
+ )
+
+ self.stage3 = nn.Sequential(
+ basic_blocks.BasicConvolutionBlock(cs[2], cs[2], ks=2, stride=2, dilation=1),
+ basic_blocks.ResidualBlock(cs[2], cs[3], ks=3, stride=1, dilation=1),
+ basic_blocks.ResidualBlock(cs[3], cs[3], ks=3, stride=1, dilation=1),
+ )
+
+ self.stage4 = nn.Sequential(
+ basic_blocks.BasicConvolutionBlock(cs[3], cs[3], ks=2, stride=2, dilation=1),
+ basic_blocks.ResidualBlock(cs[3], cs[4], ks=3, stride=1, dilation=1),
+ basic_blocks.ResidualBlock(cs[4], cs[4], ks=3, stride=1, dilation=1),
+ )
+
+ self.up1 = nn.ModuleList([
+ basic_blocks.BasicDeconvolutionBlock(cs[4], cs[5], ks=2, stride=2),
+ nn.Sequential(
+ basic_blocks.ResidualBlock(cs[5] + cs[3], cs[5], ks=3, stride=1,
+ dilation=1),
+ basic_blocks.ResidualBlock(cs[5], cs[5], ks=3, stride=1, dilation=1),
+ )
+ ])
+
+ self.up2 = nn.ModuleList([
+ basic_blocks.BasicDeconvolutionBlock(cs[5], cs[6], ks=2, stride=2),
+ nn.Sequential(
+ basic_blocks.ResidualBlock(cs[6] + cs[2], cs[6], ks=3, stride=1,
+ dilation=1),
+ basic_blocks.ResidualBlock(cs[6], cs[6], ks=3, stride=1, dilation=1),
+ )
+ ])
+
+ self.up3 = nn.ModuleList([
+ basic_blocks.BasicDeconvolutionBlock(cs[6], cs[7], ks=2, stride=2),
+ nn.Sequential(
+ basic_blocks.ResidualBlock(cs[7] + cs[1], cs[7], ks=3, stride=1,
+ dilation=1),
+ basic_blocks.ResidualBlock(cs[7], cs[7], ks=3, stride=1, dilation=1),
+ )
+ ])
+
+ self.up4 = nn.ModuleList([
+ basic_blocks.BasicDeconvolutionBlock(cs[7], cs[8], ks=2, stride=2),
+ nn.Sequential(
+ basic_blocks.ResidualBlock(cs[8] + cs[0], cs[8], ks=3, stride=1,
+ dilation=1),
+ basic_blocks.ResidualBlock(cs[8], cs[8], ks=3, stride=1, dilation=1),
+ )
+ ])
+
+ self.classifier = nn.Sequential(nn.Linear(cs[8], self.num_classes))
+
+ self.point_transforms = nn.ModuleList([
+ nn.Sequential(
+ nn.Linear(cs[0], cs[4]),
+ nn.BatchNorm1d(cs[4]),
+ nn.ReLU(True),
+ ),
+ nn.Sequential(
+ nn.Linear(cs[4], cs[6]),
+ nn.BatchNorm1d(cs[6]),
+ nn.ReLU(True),
+ ),
+ nn.Sequential(
+ nn.Linear(cs[6], cs[8]),
+ nn.BatchNorm1d(cs[8]),
+ nn.ReLU(True),
+ )
+ ])
+
+ self.weight_initialization()
+ self.dropout = nn.Dropout(0.3, True)
+
+ def weight_initialization(self):
+ for m in self.modules():
+ if isinstance(m, nn.BatchNorm1d):
+ nn.init.constant_(m.weight, 1)
+ nn.init.constant_(m.bias, 0)
+
+ def forward(self, data_dict, return_logits=False, return_final_logits=False):
+ x = data_dict['lidar']
+
+ # x: SparseTensor z: PointTensor
+ z = PointTensor(x.F, x.C.float())
+
+ x0 = initial_voxelize(z, self.pres, self.vres)
+
+ x0 = self.stem(x0)
+ z0 = voxel_to_point(x0, z, nearest=False)
+ z0.F = z0.F
+
+ x1 = point_to_voxel(x0, z0)
+ x1 = self.stage1(x1)
+ x2 = self.stage2(x1)
+ x3 = self.stage3(x2)
+ x4 = self.stage4(x3)
+ z1 = voxel_to_point(x4, z0)
+ z1.F = z1.F + self.point_transforms[0](z0.F)
+
+ y1 = point_to_voxel(x4, z1)
+
+ if return_logits:
+ output_dict = dict()
+ output_dict['logits'] = y1.F
+ output_dict['batch_indices'] = y1.C[:, -1]
+ return output_dict
+
+ y1.F = self.dropout(y1.F)
+ y1 = self.up1[0](y1)
+ y1 = torchsparse.cat([y1, x3])
+ y1 = self.up1[1](y1)
+
+ y2 = self.up2[0](y1)
+ y2 = torchsparse.cat([y2, x2])
+ y2 = self.up2[1](y2)
+ z2 = voxel_to_point(y2, z1)
+ z2.F = z2.F + self.point_transforms[1](z1.F)
+
+ y3 = point_to_voxel(y2, z2)
+ y3.F = self.dropout(y3.F)
+ y3 = self.up3[0](y3)
+ y3 = torchsparse.cat([y3, x1])
+ y3 = self.up3[1](y3)
+
+ y4 = self.up4[0](y3)
+ y4 = torchsparse.cat([y4, x0])
+ y4 = self.up4[1](y4)
+ z3 = voxel_to_point(y4, z2)
+ z3.F = z3.F + self.point_transforms[2](z2.F)
+
+ if return_final_logits:
+ output_dict = dict()
+ output_dict['logits'] = z3.F
+ output_dict['coords'] = z3.C[:, :3]
+ output_dict['batch_indices'] = z3.C[:, -1].long()
+ return output_dict
+
+ # output = self.classifier(z3.F)
+ data_dict['logits'] = z3.F
+
+ return data_dict
diff --git a/lidm/eval/models/ts/__init__.py b/lidm/eval/models/ts/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/lidm/eval/models/ts/basic_blocks.py b/lidm/eval/models/ts/basic_blocks.py
new file mode 100644
index 0000000000000000000000000000000000000000..a18acc8eba8ad5f62ad6edb2cc02852a1536d0a3
--- /dev/null
+++ b/lidm/eval/models/ts/basic_blocks.py
@@ -0,0 +1,79 @@
+#!/usr/bin/env python
+# encoding: utf-8
+'''
+@author: Xu Yan
+@file: basic_blocks.py
+@time: 2021/4/14 22:53
+'''
+import torch.nn as nn
+
+try:
+ import torchsparse.nn as spnn
+except:
+ print('To install torchsparse 1.4.0, please refer to https://github.com/mit-han-lab/torchsparse/tree/74099d10a51c71c14318bce63d6421f698b24f24')
+
+
+class BasicConvolutionBlock(nn.Module):
+ def __init__(self, inc, outc, ks=3, stride=1, dilation=1):
+ super().__init__()
+ self.net = nn.Sequential(
+ spnn.Conv3d(
+ inc,
+ outc,
+ kernel_size=ks,
+ dilation=dilation,
+ stride=stride), spnn.BatchNorm(outc),
+ spnn.ReLU(True))
+
+ def forward(self, x):
+ out = self.net(x)
+ return out
+
+
+class BasicDeconvolutionBlock(nn.Module):
+ def __init__(self, inc, outc, ks=3, stride=1):
+ super().__init__()
+ self.net = nn.Sequential(
+ spnn.Conv3d(
+ inc,
+ outc,
+ kernel_size=ks,
+ stride=stride,
+ transposed=True),
+ spnn.BatchNorm(outc),
+ spnn.ReLU(True))
+
+ def forward(self, x):
+ return self.net(x)
+
+
+class ResidualBlock(nn.Module):
+ def __init__(self, inc, outc, ks=3, stride=1, dilation=1):
+ super().__init__()
+ self.net = nn.Sequential(
+ spnn.Conv3d(
+ inc,
+ outc,
+ kernel_size=ks,
+ dilation=dilation,
+ stride=stride), spnn.BatchNorm(outc),
+ spnn.ReLU(True),
+ spnn.Conv3d(
+ outc,
+ outc,
+ kernel_size=ks,
+ dilation=dilation,
+ stride=1),
+ spnn.BatchNorm(outc))
+
+ self.downsample = nn.Sequential() if (inc == outc and stride == 1) else \
+ nn.Sequential(
+ spnn.Conv3d(inc, outc, kernel_size=1, dilation=1, stride=stride),
+ spnn.BatchNorm(outc)
+ )
+
+ self.ReLU = spnn.ReLU(True)
+
+ def forward(self, x):
+ out = self.ReLU(self.net(x) + self.downsample(x))
+ return out
diff --git a/lidm/eval/models/ts/utils.py b/lidm/eval/models/ts/utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..a4a01c0e54645ceb91fa76ed18bdae552f7fbf73
--- /dev/null
+++ b/lidm/eval/models/ts/utils.py
@@ -0,0 +1,90 @@
+import torch
+
+try:
+ import torchsparse.nn.functional as F
+ from torchsparse import PointTensor, SparseTensor
+ from torchsparse.nn.utils import get_kernel_offsets
+except:
+ print('To install torchsparse 1.4.0, please refer to https://github.com/mit-han-lab/torchsparse/tree/74099d10a51c71c14318bce63d6421f698b24f24')
+
+__all__ = ['initial_voxelize', 'point_to_voxel', 'voxel_to_point']
+
+
+# z: PointTensor
+# return: SparseTensor
+def initial_voxelize(z, init_res, after_res):
+ new_float_coord = torch.cat([(z.C[:, :3] * init_res) / after_res, z.C[:, -1].view(-1, 1)], 1)
+
+ pc_hash = F.sphash(torch.floor(new_float_coord).int())
+ sparse_hash = torch.unique(pc_hash)
+ idx_query = F.sphashquery(pc_hash, sparse_hash)
+ counts = F.spcount(idx_query.int(), len(sparse_hash))
+
+ inserted_coords = F.spvoxelize(torch.floor(new_float_coord), idx_query, counts)
+ inserted_coords = torch.round(inserted_coords).int()
+ inserted_feat = F.spvoxelize(z.F, idx_query, counts)
+
+ new_tensor = SparseTensor(inserted_feat, inserted_coords, 1)
+ new_tensor.cmaps.setdefault(new_tensor.stride, new_tensor.coords)
+ z.additional_features['idx_query'][1] = idx_query
+ z.additional_features['counts'][1] = counts
+ z.C = new_float_coord
+
+ return new_tensor
+
+
+# x: SparseTensor, z: PointTensor
+# return: SparseTensor
+def point_to_voxel(x, z):
+ if z.additional_features is None or \
+ z.additional_features.get('idx_query') is None or \
+ z.additional_features['idx_query'].get(x.s) is None:
+ pc_hash = F.sphash(
+ torch.cat([torch.floor(z.C[:, :3] / x.s[0]).int() * x.s[0], z.C[:, -1].int().view(-1, 1)], 1))
+ sparse_hash = F.sphash(x.C)
+ idx_query = F.sphashquery(pc_hash, sparse_hash)
+ counts = F.spcount(idx_query.int(), x.C.shape[0])
+ z.additional_features['idx_query'][x.s] = idx_query
+ z.additional_features['counts'][x.s] = counts
+ else:
+ idx_query = z.additional_features['idx_query'][x.s]
+ counts = z.additional_features['counts'][x.s]
+
+ inserted_feat = F.spvoxelize(z.F, idx_query, counts)
+ new_tensor = SparseTensor(inserted_feat, x.C, x.s)
+ new_tensor.cmaps = x.cmaps
+ new_tensor.kmaps = x.kmaps
+
+ return new_tensor
+
+
+# x: SparseTensor, z: PointTensor
+# return: PointTensor
+def voxel_to_point(x, z, nearest=False):
+ if z.idx_query is None or z.weights is None or z.idx_query.get(x.s) is None or z.weights.get(x.s) is None:
+ off = get_kernel_offsets(2, x.s, 1, device=z.F.device)
+ old_hash = F.sphash(
+ torch.cat([
+ torch.floor(z.C[:, :3] / x.s[0]).int() * x.s[0],
+ z.C[:, -1].int().view(-1, 1)], 1), off)
+ pc_hash = F.sphash(x.C.to(z.F.device))
+ idx_query = F.sphashquery(old_hash, pc_hash)
+ weights = F.calc_ti_weights(z.C, idx_query, scale=x.s[0]).transpose(0, 1).contiguous()
+ idx_query = idx_query.transpose(0, 1).contiguous()
+ if nearest:
+ weights[:, 1:] = 0.
+ idx_query[:, 1:] = -1
+ new_feat = F.spdevoxelize(x.F, idx_query, weights)
+ new_tensor = PointTensor(new_feat, z.C, idx_query=z.idx_query, weights=z.weights)
+ new_tensor.additional_features = z.additional_features
+ new_tensor.idx_query[x.s] = idx_query
+ new_tensor.weights[x.s] = weights
+ z.idx_query[x.s] = idx_query
+ z.weights[x.s] = weights
+
+ else:
+ new_feat = F.spdevoxelize(x.F, z.idx_query.get(x.s), z.weights.get(x.s))
+ new_tensor = PointTensor(new_feat, z.C, idx_query=z.idx_query, weights=z.weights)
+ new_tensor.additional_features = z.additional_features
+
+ return new_tensor
\ No newline at end of file
diff --git a/lidm/eval/modules/__init__.py b/lidm/eval/modules/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/lidm/eval/modules/chamfer2D/__init__.py b/lidm/eval/modules/chamfer2D/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/lidm/eval/modules/chamfer2D/chamfer2D.cu b/lidm/eval/modules/chamfer2D/chamfer2D.cu
new file mode 100755
index 0000000000000000000000000000000000000000..567dd1a0c041f0e11476e1e59bc65198ac227e04
--- /dev/null
+++ b/lidm/eval/modules/chamfer2D/chamfer2D.cu
@@ -0,0 +1,182 @@
+
+#include <stdio.h>
+#include <ATen/ATen.h>
+
+#include <cuda.h>
+#include <cuda_runtime.h>
+
+#include <vector>
+
+
+
+__global__ void NmDistanceKernel(int b,int n,const float * xyz,int m,const float * xyz2,float * result,int * result_i){
+ const int batch=512;
+ __shared__ float buf[batch*2];
+ for (int i=blockIdx.x;i<b;i+=gridDim.x){
+ for (int k2=0;k2<m;k2+=batch){
+ int end_k=min(m,k2+batch)-k2;
+ for (int j=threadIdx.x;j<end_k*2;j+=blockDim.x){
+ buf[j]=xyz2[(i*m+k2)*2+j];
+ }
+ __syncthreads();
+ for (int j=threadIdx.x+blockIdx.y*blockDim.x;j<n;j+=blockDim.x*gridDim.y){
+ float x1=xyz[(i*n+j)*2+0];
+ float y1=xyz[(i*n+j)*2+1];
+ int best_i=0;
+ float best=0;
+ int end_ka=end_k-(end_k&2);
+ if (end_ka==batch){
+ for (int k=0;k<batch;k+=4){
+ {
+ float x2=buf[k*2+0]-x1;
+ float y2=buf[k*2+1]-y1;
+ float d=x2*x2+y2*y2;
+ if (k==0 || d<best){
+ best=d;
+ best_i=k+k2;
+ }
+ }
+ {
+ float x2=buf[k*2+2]-x1;
+ float y2=buf[k*2+3]-y1;
+ float d=x2*x2+y2*y2;
+ if (d<best){
+ best=d;
+ best_i=k+k2+1;
+ }
+ }
+ {
+ float x2=buf[k*2+4]-x1;
+ float y2=buf[k*2+5]-y1;
+ float d=x2*x2+y2*y2;
+ if (d<best){
+ best=d;
+ best_i=k+k2+2;
+ }
+ }
+ {
+ float x2=buf[k*2+6]-x1;
+ float y2=buf[k*2+7]-y1;
+ float d=x2*x2+y2*y2;
+ if (d<best){
+ best=d;
+ best_i=k+k2+3;
+ }
+ }
+ }
+ }else{
+ for (int k=0;k<end_ka;k+=4){
+ {
+ float x2=buf[k*2+0]-x1;
+ float y2=buf[k*2+1]-y1;
+ float d=x2*x2+y2*y2;
+ if (k==0 || d<best){
+ best=d;
+ best_i=k+k2;
+ }
+ }
+ {
+ float x2=buf[k*2+2]-x1;
+ float y2=buf[k*2+3]-y1;
+ float d=x2*x2+y2*y2;
+ if (d<best){
+ best=d;
+ best_i=k+k2+1;
+ }
+ }
+ {
+ float x2=buf[k*2+4]-x1;
+ float y2=buf[k*2+5]-y1;
+ float d=x2*x2+y2*y2;
+ if (d<best){
+ best=d;
+ best_i=k+k2+2;
+ }
+ }
+ {
+ float x2=buf[k*2+6]-x1;
+ float y2=buf[k*2+7]-y1;
+ float d=x2*x2+y2*y2;
+ if (d<best){
+ best=d;
+ best_i=k+k2+3;
+ }
+ }
+ }
+ }
+ for (int k=end_ka;k<end_k;k++){
+ float x2=buf[k*2+0]-x1;
+ float y2=buf[k*2+1]-y1;
+ float d=x2*x2+y2*y2;
+ if (k==0 || d<best){
+ best=d;
+ best_i=k+k2;
+ }
+ }
+ if (k2==0 || result[(i*n+j)]>best){
+ result[(i*n+j)]=best;
+ result_i[(i*n+j)]=best_i;
+ }
+ }
+ __syncthreads();
+ }
+ }
+}
+// int chamfer_cuda_forward(int b,int n,const float * xyz,int m,const float * xyz2,float * result,int * result_i,float * result2,int * result2_i, cudaStream_t stream){
+int chamfer_cuda_forward(at::Tensor xyz1, at::Tensor xyz2, at::Tensor dist1, at::Tensor dist2, at::Tensor idx1, at::Tensor idx2){
+
+ const auto batch_size = xyz1.size(0);
+ const auto n = xyz1.size(1); //num_points point cloud A
+ const auto m = xyz2.size(1); //num_points point cloud B
+
+ NmDistanceKernel<<<dim3(32,16,1),512>>>(batch_size, n, xyz1.data<float>(), m, xyz2.data<float>(), dist1.data<float>(), idx1.data<int>());
+ NmDistanceKernel<<<dim3(32,16,1),512>>>(batch_size, m, xyz2.data<float>(), n, xyz1.data<float>(), dist2.data<float>(), idx2.data<int>());
+
+ cudaError_t err = cudaGetLastError();
+ if (err != cudaSuccess) {
+ printf("error in nnd updateOutput: %s\n", cudaGetErrorString(err));
+ //THError("aborting");
+ return 0;
+ }
+ return 1;
+
+
+}
+__global__ void NmDistanceGradKernel(int b,int n,const float * xyz1,int m,const float * xyz2,const float * grad_dist1,const int * idx1,float * grad_xyz1,float * grad_xyz2){
+ for (int i=blockIdx.x;i<b;i+=gridDim.x){
+ for (int j=threadIdx.x+blockIdx.y*blockDim.x;j<n;j+=blockDim.x*gridDim.y){
+ float x1=xyz1[(i*n+j)*2+0];
+ float y1=xyz1[(i*n+j)*2+1];
+ int j2=idx1[i*n+j];
+ float x2=xyz2[(i*m+j2)*2+0];
+ float y2=xyz2[(i*m+j2)*2+1];
+ float g=grad_dist1[i*n+j]*2;
+ atomicAdd(&(grad_xyz1[(i*n+j)*2+0]),g*(x1-x2));
+ atomicAdd(&(grad_xyz1[(i*n+j)*2+1]),g*(y1-y2));
+ atomicAdd(&(grad_xyz2[(i*m+j2)*2+0]),-(g*(x1-x2)));
+ atomicAdd(&(grad_xyz2[(i*m+j2)*2+1]),-(g*(y1-y2)));
+ }
+ }
+}
+// int chamfer_cuda_backward(int b,int n,const float * xyz1,int m,const float * xyz2,const float * grad_dist1,const int * idx1,const float * grad_dist2,const int * idx2,float * grad_xyz1,float * grad_xyz2, cudaStream_t stream){
+int chamfer_cuda_backward(at::Tensor xyz1, at::Tensor xyz2, at::Tensor gradxyz1, at::Tensor gradxyz2, at::Tensor graddist1, at::Tensor graddist2, at::Tensor idx1, at::Tensor idx2){
+ // cudaMemset(grad_xyz1,0,b*n*3*4);
+ // cudaMemset(grad_xyz2,0,b*m*3*4);
+
+ const auto batch_size = xyz1.size(0);
+ const auto n = xyz1.size(1); //num_points point cloud A
+ const auto m = xyz2.size(1); //num_points point cloud B
+
+ NmDistanceGradKernel<<<dim3(1,16,1),256>>>(batch_size,n,xyz1.data<float>(),m,xyz2.data<float>(),graddist1.data<float>(),idx1.data<int>(),gradxyz1.data<float>(),gradxyz2.data<float>());
+ NmDistanceGradKernel<<<dim3(1,16,1),256>>>(batch_size,m,xyz2.data<float>(),n,xyz1.data<float>(),graddist2.data<float>(),idx2.data<int>(),gradxyz2.data<float>(),gradxyz1.data<float>());
+
+ cudaError_t err = cudaGetLastError();
+ if (err != cudaSuccess) {
+ printf("error in nnd get grad: %s\n", cudaGetErrorString(err));
+ //THError("aborting");
+ return 0;
+ }
+ return 1;
+
+}
+
diff --git a/lidm/eval/modules/chamfer2D/chamfer_cuda.cpp b/lidm/eval/modules/chamfer2D/chamfer_cuda.cpp
new file mode 100755
index 0000000000000000000000000000000000000000..67574e21818ae9388f44f964d606b2f00355865e
--- /dev/null
+++ b/lidm/eval/modules/chamfer2D/chamfer_cuda.cpp
@@ -0,0 +1,33 @@
+#include <torch/torch.h>
+#include <vector>
+
+///TMP
+//#include "common.h"
+/// NOT TMP
+
+
+int chamfer_cuda_forward(at::Tensor xyz1, at::Tensor xyz2, at::Tensor dist1, at::Tensor dist2, at::Tensor idx1, at::Tensor idx2);
+
+
+int chamfer_cuda_backward(at::Tensor xyz1, at::Tensor xyz2, at::Tensor gradxyz1, at::Tensor gradxyz2, at::Tensor graddist1, at::Tensor graddist2, at::Tensor idx1, at::Tensor idx2);
+
+
+
+
+int chamfer_forward(at::Tensor xyz1, at::Tensor xyz2, at::Tensor dist1, at::Tensor dist2, at::Tensor idx1, at::Tensor idx2) {
+ return chamfer_cuda_forward(xyz1, xyz2, dist1, dist2, idx1, idx2);
+}
+
+
+int chamfer_backward(at::Tensor xyz1, at::Tensor xyz2, at::Tensor gradxyz1, at::Tensor gradxyz2, at::Tensor graddist1,
+ at::Tensor graddist2, at::Tensor idx1, at::Tensor idx2) {
+
+ return chamfer_cuda_backward(xyz1, xyz2, gradxyz1, gradxyz2, graddist1, graddist2, idx1, idx2);
+}
+
+
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+ m.def("forward", &chamfer_forward, "chamfer forward (CUDA)");
+ m.def("backward", &chamfer_backward, "chamfer backward (CUDA)");
+}
\ No newline at end of file
diff --git a/lidm/eval/modules/chamfer2D/dist_chamfer_2D.py b/lidm/eval/modules/chamfer2D/dist_chamfer_2D.py
new file mode 100644
index 0000000000000000000000000000000000000000..e7ecedbc731b23895ed2d6fe3fe48aaf632940da
--- /dev/null
+++ b/lidm/eval/modules/chamfer2D/dist_chamfer_2D.py
@@ -0,0 +1,84 @@
+from torch import nn
+from torch.autograd import Function
+import torch
+import importlib
+import os
+
+chamfer_found = importlib.find_loader("chamfer_2D") is not None
+if not chamfer_found:
+ ## Cool trick from https://github.com/chrdiller
+ print("Jitting Chamfer 2D")
+ cur_path = os.path.dirname(os.path.abspath(__file__))
+ build_path = cur_path.replace('chamfer2D', 'tmp')
+ os.makedirs(build_path, exist_ok=True)
+
+ from torch.utils.cpp_extension import load
+
+ chamfer_2D = load(name="chamfer_2D",
+ sources=[
+ "/".join(os.path.abspath(__file__).split('/')[:-1] + ["chamfer_cuda.cpp"]),
+ "/".join(os.path.abspath(__file__).split('/')[:-1] + ["chamfer2D.cu"]),
+ ], build_directory=build_path)
+ print("Loaded JIT 2D CUDA chamfer distance")
+
+else:
+ import chamfer_2D
+
+ print("Loaded compiled 2D CUDA chamfer distance")
+
+
+# Chamfer's distance module @thibaultgroueix
+# GPU tensors only
+class chamfer_2DFunction(Function):
+ @staticmethod
+ def forward(ctx, xyz1, xyz2):
+ batchsize, n, dim = xyz1.size()
+ assert dim == 2, "Wrong last dimension for the chamfer distance 's input! Check with .size()"
+ _, m, dim = xyz2.size()
+ assert dim == 2, "Wrong last dimension for the chamfer distance 's input! Check with .size()"
+ device = xyz1.device
+
+ device = xyz1.device
+
+ dist1 = torch.zeros(batchsize, n)
+ dist2 = torch.zeros(batchsize, m)
+
+ idx1 = torch.zeros(batchsize, n).type(torch.IntTensor)
+ idx2 = torch.zeros(batchsize, m).type(torch.IntTensor)
+
+ dist1 = dist1.to(device)
+ dist2 = dist2.to(device)
+ idx1 = idx1.to(device)
+ idx2 = idx2.to(device)
+ torch.cuda.set_device(device)
+
+ chamfer_2D.forward(xyz1, xyz2, dist1, dist2, idx1, idx2)
+ ctx.save_for_backward(xyz1, xyz2, idx1, idx2)
+ return dist1, dist2, idx1, idx2
+
+ @staticmethod
+ def backward(ctx, graddist1, graddist2, gradidx1, gradidx2):
+ xyz1, xyz2, idx1, idx2 = ctx.saved_tensors
+ graddist1 = graddist1.contiguous()
+ graddist2 = graddist2.contiguous()
+ device = graddist1.device
+
+ gradxyz1 = torch.zeros(xyz1.size())
+ gradxyz2 = torch.zeros(xyz2.size())
+
+ gradxyz1 = gradxyz1.to(device)
+ gradxyz2 = gradxyz2.to(device)
+ chamfer_2D.backward(
+ xyz1, xyz2, gradxyz1, gradxyz2, graddist1, graddist2, idx1, idx2
+ )
+ return gradxyz1, gradxyz2
+
+
+class chamfer_2DDist(nn.Module):
+ def __init__(self):
+ super(chamfer_2DDist, self).__init__()
+
+ def forward(self, input1, input2):
+ input1 = input1.contiguous()
+ input2 = input2.contiguous()
+ return chamfer_2DFunction.apply(input1, input2)
diff --git a/lidm/eval/modules/chamfer2D/setup.py b/lidm/eval/modules/chamfer2D/setup.py
new file mode 100755
index 0000000000000000000000000000000000000000..11d01237f57ad386dd88adda4cc869a53f94f4f2
--- /dev/null
+++ b/lidm/eval/modules/chamfer2D/setup.py
@@ -0,0 +1,14 @@
+from setuptools import setup
+from torch.utils.cpp_extension import BuildExtension, CUDAExtension
+
+setup(
+ name='chamfer_2D',
+ ext_modules=[
+ CUDAExtension('chamfer_2D', [
+ "/".join(__file__.split('/')[:-1] + ['chamfer_cuda.cpp']),
+ "/".join(__file__.split('/')[:-1] + ['chamfer2D.cu']),
+ ]),
+ ],
+ cmdclass={
+ 'build_ext': BuildExtension
+ })
\ No newline at end of file
diff --git a/lidm/eval/modules/chamfer3D/__init__.py b/lidm/eval/modules/chamfer3D/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/lidm/eval/modules/chamfer3D/chamfer3D.cu b/lidm/eval/modules/chamfer3D/chamfer3D.cu
new file mode 100644
index 0000000000000000000000000000000000000000..d5b886dff11733be30519247d1fdb784818bff4a
--- /dev/null
+++ b/lidm/eval/modules/chamfer3D/chamfer3D.cu
@@ -0,0 +1,196 @@
+
+#include <stdio.h>
+#include <ATen/ATen.h>
+
+#include <cuda.h>
+#include <cuda_runtime.h>
+
+#include <vector>
+
+
+
+__global__ void NmDistanceKernel(int b,int n,const float * xyz,int m,const float * xyz2,float * result,int * result_i){
+ const int batch=512;
+ __shared__ float buf[batch*3];
+ for (int i=blockIdx.x;i<b;i+=gridDim.x){
+ for (int k2=0;k2<m;k2+=batch){
+ int end_k=min(m,k2+batch)-k2;
+ for (int j=threadIdx.x;j<end_k*3;j+=blockDim.x){
+ buf[j]=xyz2[(i*m+k2)*3+j];
+ }
+ __syncthreads();
+ for (int j=threadIdx.x+blockIdx.y*blockDim.x;j<n;j+=blockDim.x*gridDim.y){
+ float x1=xyz[(i*n+j)*3+0];
+ float y1=xyz[(i*n+j)*3+1];
+ float z1=xyz[(i*n+j)*3+2];
+ int best_i=0;
+ float best=0;
+ int end_ka=end_k-(end_k&3);
+ if (end_ka==batch){
+ for (int k=0;k<batch;k+=4){
+ {
+ float x2=buf[k*3+0]-x1;
+ float y2=buf[k*3+1]-y1;
+ float z2=buf[k*3+2]-z1;
+ float d=x2*x2+y2*y2+z2*z2;
+ if (k==0 || d<best){
+ best=d;
+ best_i=k+k2;
+ }
+ }
+ {
+ float x2=buf[k*3+3]-x1;
+ float y2=buf[k*3+4]-y1;
+ float z2=buf[k*3+5]-z1;
+ float d=x2*x2+y2*y2+z2*z2;
+ if (d<best){
+ best=d;
+ best_i=k+k2+1;
+ }
+ }
+ {
+ float x2=buf[k*3+6]-x1;
+ float y2=buf[k*3+7]-y1;
+ float z2=buf[k*3+8]-z1;
+ float d=x2*x2+y2*y2+z2*z2;
+ if (d<best){
+ best=d;
+ best_i=k+k2+2;
+ }
+ }
+ {
+ float x2=buf[k*3+9]-x1;
+ float y2=buf[k*3+10]-y1;
+ float z2=buf[k*3+11]-z1;
+ float d=x2*x2+y2*y2+z2*z2;
+ if (d<best){
+ best=d;
+ best_i=k+k2+3;
+ }
+ }
+ }
+ }else{
+ for (int k=0;k<end_ka;k+=4){
+ {
+ float x2=buf[k*3+0]-x1;
+ float y2=buf[k*3+1]-y1;
+ float z2=buf[k*3+2]-z1;
+ float d=x2*x2+y2*y2+z2*z2;
+ if (k==0 || d<best){
+ best=d;
+ best_i=k+k2;
+ }
+ }
+ {
+ float x2=buf[k*3+3]-x1;
+ float y2=buf[k*3+4]-y1;
+ float z2=buf[k*3+5]-z1;
+ float d=x2*x2+y2*y2+z2*z2;
+ if (d<best){
+ best=d;
+ best_i=k+k2+1;
+ }
+ }
+ {
+ float x2=buf[k*3+6]-x1;
+ float y2=buf[k*3+7]-y1;
+ float z2=buf[k*3+8]-z1;
+ float d=x2*x2+y2*y2+z2*z2;
+ if (d<best){
+ best=d;
+ best_i=k+k2+2;
+ }
+ }
+ {
+ float x2=buf[k*3+9]-x1;
+ float y2=buf[k*3+10]-y1;
+ float z2=buf[k*3+11]-z1;
+ float d=x2*x2+y2*y2+z2*z2;
+ if (d<best){
+ best=d;
+ best_i=k+k2+3;
+ }
+ }
+ }
+ }
+ for (int k=end_ka;k<end_k;k++){
+ float x2=buf[k*3+0]-x1;
+ float y2=buf[k*3+1]-y1;
+ float z2=buf[k*3+2]-z1;
+ float d=x2*x2+y2*y2+z2*z2;
+ if (k==0 || d<best){
+ best=d;
+ best_i=k+k2;
+ }
+ }
+ if (k2==0 || result[(i*n+j)]>best){
+ result[(i*n+j)]=best;
+ result_i[(i*n+j)]=best_i;
+ }
+ }
+ __syncthreads();
+ }
+ }
+}
+// int chamfer_cuda_forward(int b,int n,const float * xyz,int m,const float * xyz2,float * result,int * result_i,float * result2,int * result2_i, cudaStream_t stream){
+int chamfer_cuda_forward(at::Tensor xyz1, at::Tensor xyz2, at::Tensor dist1, at::Tensor dist2, at::Tensor idx1, at::Tensor idx2){
+
+ const auto batch_size = xyz1.size(0);
+ const auto n = xyz1.size(1); //num_points point cloud A
+ const auto m = xyz2.size(1); //num_points point cloud B
+
+ NmDistanceKernel<<<dim3(32,16,1),512>>>(batch_size, n, xyz1.data<float>(), m, xyz2.data<float>(), dist1.data<float>(), idx1.data<int>());
+ NmDistanceKernel<<<dim3(32,16,1),512>>>(batch_size, m, xyz2.data<float>(), n, xyz1.data<float>(), dist2.data<float>(), idx2.data<int>());
+
+ cudaError_t err = cudaGetLastError();
+ if (err != cudaSuccess) {
+ printf("error in nnd updateOutput: %s\n", cudaGetErrorString(err));
+ //THError("aborting");
+ return 0;
+ }
+ return 1;
+
+
+}
+__global__ void NmDistanceGradKernel(int b,int n,const float * xyz1,int m,const float * xyz2,const float * grad_dist1,const int * idx1,float * grad_xyz1,float * grad_xyz2){
+ for (int i=blockIdx.x;i<b;i+=gridDim.x){
+ for (int j=threadIdx.x+blockIdx.y*blockDim.x;j<n;j+=blockDim.x*gridDim.y){
+ float x1=xyz1[(i*n+j)*3+0];
+ float y1=xyz1[(i*n+j)*3+1];
+ float z1=xyz1[(i*n+j)*3+2];
+ int j2=idx1[i*n+j];
+ float x2=xyz2[(i*m+j2)*3+0];
+ float y2=xyz2[(i*m+j2)*3+1];
+ float z2=xyz2[(i*m+j2)*3+2];
+ float g=grad_dist1[i*n+j]*2;
+ atomicAdd(&(grad_xyz1[(i*n+j)*3+0]),g*(x1-x2));
+ atomicAdd(&(grad_xyz1[(i*n+j)*3+1]),g*(y1-y2));
+ atomicAdd(&(grad_xyz1[(i*n+j)*3+2]),g*(z1-z2));
+ atomicAdd(&(grad_xyz2[(i*m+j2)*3+0]),-(g*(x1-x2)));
+ atomicAdd(&(grad_xyz2[(i*m+j2)*3+1]),-(g*(y1-y2)));
+ atomicAdd(&(grad_xyz2[(i*m+j2)*3+2]),-(g*(z1-z2)));
+ }
+ }
+}
+// int chamfer_cuda_backward(int b,int n,const float * xyz1,int m,const float * xyz2,const float * grad_dist1,const int * idx1,const float * grad_dist2,const int * idx2,float * grad_xyz1,float * grad_xyz2, cudaStream_t stream){
+int chamfer_cuda_backward(at::Tensor xyz1, at::Tensor xyz2, at::Tensor gradxyz1, at::Tensor gradxyz2, at::Tensor graddist1, at::Tensor graddist2, at::Tensor idx1, at::Tensor idx2){
+ // cudaMemset(grad_xyz1,0,b*n*3*4);
+ // cudaMemset(grad_xyz2,0,b*m*3*4);
+
+ const auto batch_size = xyz1.size(0);
+ const auto n = xyz1.size(1); //num_points point cloud A
+ const auto m = xyz2.size(1); //num_points point cloud B
+
+ NmDistanceGradKernel<<<dim3(1,16,1),256>>>(batch_size,n,xyz1.data<float>(),m,xyz2.data<float>(),graddist1.data<float>(),idx1.data<int>(),gradxyz1.data<float>(),gradxyz2.data<float>());
+ NmDistanceGradKernel<<<dim3(1,16,1),256>>>(batch_size,m,xyz2.data<float>(),n,xyz1.data<float>(),graddist2.data<float>(),idx2.data<int>(),gradxyz2.data<float>(),gradxyz1.data<float>());
+
+ cudaError_t err = cudaGetLastError();
+ if (err != cudaSuccess) {
+ printf("error in nnd get grad: %s\n", cudaGetErrorString(err));
+ //THError("aborting");
+ return 0;
+ }
+ return 1;
+
+}
+
diff --git a/lidm/eval/modules/chamfer3D/chamfer_cuda.cpp b/lidm/eval/modules/chamfer3D/chamfer_cuda.cpp
new file mode 100644
index 0000000000000000000000000000000000000000..67574e21818ae9388f44f964d606b2f00355865e
--- /dev/null
+++ b/lidm/eval/modules/chamfer3D/chamfer_cuda.cpp
@@ -0,0 +1,33 @@
+#include <torch/torch.h>
+#include <vector>
+
+///TMP
+//#include "common.h"
+/// NOT TMP
+
+
+int chamfer_cuda_forward(at::Tensor xyz1, at::Tensor xyz2, at::Tensor dist1, at::Tensor dist2, at::Tensor idx1, at::Tensor idx2);
+
+
+int chamfer_cuda_backward(at::Tensor xyz1, at::Tensor xyz2, at::Tensor gradxyz1, at::Tensor gradxyz2, at::Tensor graddist1, at::Tensor graddist2, at::Tensor idx1, at::Tensor idx2);
+
+
+
+
+int chamfer_forward(at::Tensor xyz1, at::Tensor xyz2, at::Tensor dist1, at::Tensor dist2, at::Tensor idx1, at::Tensor idx2) {
+ return chamfer_cuda_forward(xyz1, xyz2, dist1, dist2, idx1, idx2);
+}
+
+
+int chamfer_backward(at::Tensor xyz1, at::Tensor xyz2, at::Tensor gradxyz1, at::Tensor gradxyz2, at::Tensor graddist1,
+ at::Tensor graddist2, at::Tensor idx1, at::Tensor idx2) {
+
+ return chamfer_cuda_backward(xyz1, xyz2, gradxyz1, gradxyz2, graddist1, graddist2, idx1, idx2);
+}
+
+
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+ m.def("forward", &chamfer_forward, "chamfer forward (CUDA)");
+ m.def("backward", &chamfer_backward, "chamfer backward (CUDA)");
+}
\ No newline at end of file
diff --git a/lidm/eval/modules/chamfer3D/dist_chamfer_3D.py b/lidm/eval/modules/chamfer3D/dist_chamfer_3D.py
new file mode 100644
index 0000000000000000000000000000000000000000..30063e50293445f0785f1cb7fd05719a1a29b172
--- /dev/null
+++ b/lidm/eval/modules/chamfer3D/dist_chamfer_3D.py
@@ -0,0 +1,76 @@
+from torch import nn
+from torch.autograd import Function
+import torch
+import importlib
+import os
+
+chamfer_found = importlib.find_loader("chamfer_3D") is not None
+if not chamfer_found:
+ ## Cool trick from https://github.com/chrdiller
+ print("Jitting Chamfer 3D")
+
+ from torch.utils.cpp_extension import load
+
+ chamfer_3D = load(name="chamfer_3D",
+ sources=[
+ "/".join(os.path.abspath(__file__).split('/')[:-1] + ["chamfer_cuda.cpp"]),
+ "/".join(os.path.abspath(__file__).split('/')[:-1] + ["chamfer3D.cu"]),
+ ])
+ print("Loaded JIT 3D CUDA chamfer distance")
+
+else:
+ import chamfer_3D
+ print("Loaded compiled 3D CUDA chamfer distance")
+
+
+# Chamfer's distance module @thibaultgroueix
+# GPU tensors only
+class chamfer_3DFunction(Function):
+ @staticmethod
+ def forward(ctx, xyz1, xyz2):
+ batchsize, n, _ = xyz1.size()
+ _, m, _ = xyz2.size()
+ device = xyz1.device
+
+ dist1 = torch.zeros(batchsize, n)
+ dist2 = torch.zeros(batchsize, m)
+
+ idx1 = torch.zeros(batchsize, n).type(torch.IntTensor)
+ idx2 = torch.zeros(batchsize, m).type(torch.IntTensor)
+
+ dist1 = dist1.to(device)
+ dist2 = dist2.to(device)
+ idx1 = idx1.to(device)
+ idx2 = idx2.to(device)
+ torch.cuda.set_device(device)
+
+ chamfer_3D.forward(xyz1, xyz2, dist1, dist2, idx1, idx2)
+ ctx.save_for_backward(xyz1, xyz2, idx1, idx2)
+ return dist1, dist2, idx1, idx2
+
+ @staticmethod
+ def backward(ctx, graddist1, graddist2, gradidx1, gradidx2):
+ xyz1, xyz2, idx1, idx2 = ctx.saved_tensors
+ graddist1 = graddist1.contiguous()
+ graddist2 = graddist2.contiguous()
+ device = graddist1.device
+
+ gradxyz1 = torch.zeros(xyz1.size())
+ gradxyz2 = torch.zeros(xyz2.size())
+
+ gradxyz1 = gradxyz1.to(device)
+ gradxyz2 = gradxyz2.to(device)
+ chamfer_3D.backward(
+ xyz1, xyz2, gradxyz1, gradxyz2, graddist1, graddist2, idx1, idx2
+ )
+ return gradxyz1, gradxyz2
+
+
+class chamfer_3DDist(nn.Module):
+ def __init__(self):
+ super(chamfer_3DDist, self).__init__()
+
+ def forward(self, input1, input2):
+ input1 = input1.contiguous()
+ input2 = input2.contiguous()
+ return chamfer_3DFunction.apply(input1, input2)
diff --git a/lidm/eval/modules/chamfer3D/setup.py b/lidm/eval/modules/chamfer3D/setup.py
new file mode 100644
index 0000000000000000000000000000000000000000..9a23aadadde026eb8c3db68a43d63086f6be856a
--- /dev/null
+++ b/lidm/eval/modules/chamfer3D/setup.py
@@ -0,0 +1,14 @@
+from setuptools import setup
+from torch.utils.cpp_extension import BuildExtension, CUDAExtension
+
+setup(
+ name='chamfer_3D',
+ ext_modules=[
+ CUDAExtension('chamfer_3D', [
+ "/".join(__file__.split('/')[:-1] + ['chamfer_cuda.cpp']),
+ "/".join(__file__.split('/')[:-1] + ['chamfer3D.cu']),
+ ]),
+ ],
+ cmdclass={
+ 'build_ext': BuildExtension
+ })
\ No newline at end of file
diff --git a/lidm/eval/modules/emd/__init__.py b/lidm/eval/modules/emd/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/lidm/eval/modules/emd/emd.cpp b/lidm/eval/modules/emd/emd.cpp
new file mode 100644
index 0000000000000000000000000000000000000000..5036016370aa5f148b17ec44ea7a30bc576c3a82
--- /dev/null
+++ b/lidm/eval/modules/emd/emd.cpp
@@ -0,0 +1,31 @@
+// EMD approximation module (based on auction algorithm)
+// author: Minghua Liu
+#include <torch/extension.h>
+#include <vector>
+
+int emd_cuda_forward(at::Tensor xyz1, at::Tensor xyz2, at::Tensor dist, at::Tensor assignment, at::Tensor price,
+ at::Tensor assignment_inv, at::Tensor bid, at::Tensor bid_increments, at::Tensor max_increments,
+ at::Tensor unass_idx, at::Tensor unass_cnt, at::Tensor unass_cnt_sum, at::Tensor cnt_tmp, at::Tensor max_idx, float eps, int iters);
+
+int emd_cuda_backward(at::Tensor xyz1, at::Tensor xyz2, at::Tensor gradxyz, at::Tensor graddist, at::Tensor idx);
+
+
+
+int emd_forward(at::Tensor xyz1, at::Tensor xyz2, at::Tensor dist, at::Tensor assignment, at::Tensor price,
+ at::Tensor assignment_inv, at::Tensor bid, at::Tensor bid_increments, at::Tensor max_increments,
+ at::Tensor unass_idx, at::Tensor unass_cnt, at::Tensor unass_cnt_sum, at::Tensor cnt_tmp, at::Tensor max_idx, float eps, int iters) {
+ return emd_cuda_forward(xyz1, xyz2, dist, assignment, price, assignment_inv, bid, bid_increments, max_increments, unass_idx, unass_cnt, unass_cnt_sum, cnt_tmp, max_idx, eps, iters);
+}
+
+int emd_backward(at::Tensor xyz1, at::Tensor xyz2, at::Tensor gradxyz, at::Tensor graddist, at::Tensor idx) {
+
+ return emd_cuda_backward(xyz1, xyz2, gradxyz, graddist, idx);
+}
+
+
+
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+ m.def("forward", &emd_forward, "emd forward (CUDA)");
+ m.def("backward", &emd_backward, "emd backward (CUDA)");
+}
\ No newline at end of file
diff --git a/lidm/eval/modules/emd/emd_cuda.cu b/lidm/eval/modules/emd/emd_cuda.cu
new file mode 100644
index 0000000000000000000000000000000000000000..08999bbac057bd09ed7668b71e2f691e4f06fd6b
--- /dev/null
+++ b/lidm/eval/modules/emd/emd_cuda.cu
@@ -0,0 +1,316 @@
+// EMD approximation module (based on auction algorithm)
+// author: Minghua Liu
+#include <stdio.h>
+#include <ATen/ATen.h>
+
+#include <cuda.h>
+#include <iostream>
+#include <cuda_runtime.h>
+
+__device__ __forceinline__ float atomicMax(float *address, float val)
+{
+ int ret = __float_as_int(*address);
+ while(val > __int_as_float(ret))
+ {
+ int old = ret;
+ if((ret = atomicCAS((int *)address, old, __float_as_int(val))) == old)
+ break;
+ }
+ return __int_as_float(ret);
+}
+
+
+__global__ void clear(int b, int * cnt_tmp, int * unass_cnt) {
+ for (int i = threadIdx.x; i < b; i += blockDim.x) {
+ cnt_tmp[i] = 0;
+ unass_cnt[i] = 0;
+ }
+}
+
+__global__ void calc_unass_cnt(int b, int n, int * assignment, int * unass_cnt) {
+ // count the number of unassigned points in each batch
+ const int BLOCK_SIZE = 1024;
+ __shared__ int scan_array[BLOCK_SIZE];
+ for (int i = blockIdx.x; i < b; i += gridDim.x) {
+ scan_array[threadIdx.x] = assignment[i * n + blockIdx.y * BLOCK_SIZE + threadIdx.x] == -1 ? 1 : 0;
+ __syncthreads();
+
+ int stride = 1;
+ while(stride <= BLOCK_SIZE / 2) {
+ int index = (threadIdx.x + 1) * stride * 2 - 1;
+ if(index < BLOCK_SIZE)
+ scan_array[index] += scan_array[index - stride];
+ stride = stride * 2;
+ __syncthreads();
+ }
+ __syncthreads();
+
+ if (threadIdx.x == BLOCK_SIZE - 1) {
+ atomicAdd(&unass_cnt[i], scan_array[threadIdx.x]);
+ }
+ __syncthreads();
+ }
+}
+
+__global__ void calc_unass_cnt_sum(int b, int * unass_cnt, int * unass_cnt_sum) {
+ // count the cumulative sum over over unass_cnt
+ const int BLOCK_SIZE = 512; // batch_size <= 512
+ __shared__ int scan_array[BLOCK_SIZE];
+ scan_array[threadIdx.x] = unass_cnt[threadIdx.x];
+ __syncthreads();
+
+ int stride = 1;
+ while(stride <= BLOCK_SIZE / 2) {
+ int index = (threadIdx.x + 1) * stride * 2 - 1;
+ if(index < BLOCK_SIZE)
+ scan_array[index] += scan_array[index - stride];
+ stride = stride * 2;
+ __syncthreads();
+ }
+ __syncthreads();
+ stride = BLOCK_SIZE / 4;
+ while(stride > 0) {
+ int index = (threadIdx.x + 1) * stride * 2 - 1;
+ if((index + stride) < BLOCK_SIZE)
+ scan_array[index + stride] += scan_array[index];
+ stride = stride / 2;
+ __syncthreads();
+ }
+ __syncthreads();
+
+ //printf("%d\n", unass_cnt_sum[b - 1]);
+ unass_cnt_sum[threadIdx.x] = scan_array[threadIdx.x];
+}
+
+__global__ void calc_unass_idx(int b, int n, int * assignment, int * unass_idx, int * unass_cnt, int * unass_cnt_sum, int * cnt_tmp) {
+ // list all the unassigned points
+ for (int i = blockIdx.x; i < b; i += gridDim.x) {
+ if (assignment[i * n + blockIdx.y * 1024 + threadIdx.x] == -1) {
+ int idx = atomicAdd(&cnt_tmp[i], 1);
+ unass_idx[unass_cnt_sum[i] - unass_cnt[i] + idx] = blockIdx.y * 1024 + threadIdx.x;
+ }
+ }
+}
+
+__global__ void Bid(int b, int n, const float * xyz1, const float * xyz2, float eps, int * assignment, int * assignment_inv, float * price,
+ int * bid, float * bid_increments, float * max_increments, int * unass_cnt, int * unass_cnt_sum, int * unass_idx) {
+ const int batch = 2048, block_size = 1024, block_cnt = n / 1024;
+ __shared__ float xyz2_buf[batch * 3];
+ __shared__ float price_buf[batch];
+ __shared__ float best_buf[block_size];
+ __shared__ float better_buf[block_size];
+ __shared__ int best_i_buf[block_size];
+ for (int i = blockIdx.x; i < b; i += gridDim.x) {
+ int _unass_cnt = unass_cnt[i];
+ if (_unass_cnt == 0)
+ continue;
+ int _unass_cnt_sum = unass_cnt_sum[i];
+ int unass_per_block = (_unass_cnt + block_cnt - 1) / block_cnt;
+ int thread_per_unass = block_size / unass_per_block;
+ int unass_this_block = max(min(_unass_cnt - (int) blockIdx.y * unass_per_block, unass_per_block), 0);
+
+ float x1, y1, z1, best = -1e9, better = -1e9;
+ int best_i = -1, _unass_id = -1, thread_in_unass;
+
+ if (threadIdx.x < thread_per_unass * unass_this_block) {
+ _unass_id = unass_per_block * blockIdx.y + threadIdx.x / thread_per_unass + _unass_cnt_sum - _unass_cnt;
+ _unass_id = unass_idx[_unass_id];
+ thread_in_unass = threadIdx.x % thread_per_unass;
+
+ x1 = xyz1[(i * n + _unass_id) * 3 + 0];
+ y1 = xyz1[(i * n + _unass_id) * 3 + 1];
+ z1 = xyz1[(i * n + _unass_id) * 3 + 2];
+ }
+
+ for (int k2 = 0; k2 < n; k2 += batch) {
+ int end_k = min(n, k2 + batch) - k2;
+ for (int j = threadIdx.x; j < end_k * 3; j += blockDim.x) {
+ xyz2_buf[j] = xyz2[(i * n + k2) * 3 + j];
+ }
+ for (int j = threadIdx.x; j < end_k; j += blockDim.x) {
+ price_buf[j] = price[i * n + k2 + j];
+ }
+ __syncthreads();
+
+ if (_unass_id != -1) {
+ int delta = (end_k + thread_per_unass - 1) / thread_per_unass;
+ int l = thread_in_unass * delta;
+ int r = min((thread_in_unass + 1) * delta, end_k);
+ for (int k = l; k < r; k++)
+ //if (!last || assignment_inv[i * n + k + k2] == -1)
+ {
+ float x2 = xyz2_buf[k * 3 + 0] - x1;
+ float y2 = xyz2_buf[k * 3 + 1] - y1;
+ float z2 = xyz2_buf[k * 3 + 2] - z1;
+ // the coordinates of points should be normalized to [0, 1]
+ float d = 3.0 - sqrtf(x2 * x2 + y2 * y2 + z2 * z2) - price_buf[k];
+ if (d > best) {
+ better = best;
+ best = d;
+ best_i = k + k2;
+ }
+ else if (d > better) {
+ better = d;
+ }
+ }
+ }
+ __syncthreads();
+ }
+
+ best_buf[threadIdx.x] = best;
+ better_buf[threadIdx.x] = better;
+ best_i_buf[threadIdx.x] = best_i;
+ __syncthreads();
+
+ if (_unass_id != -1 && thread_in_unass == 0) {
+ for (int j = threadIdx.x + 1; j < threadIdx.x + thread_per_unass; j++) {
+ if (best_buf[j] > best) {
+ better = max(best, better_buf[j]);
+ best = best_buf[j];
+ best_i = best_i_buf[j];
+ }
+ else better = max(better, best_buf[j]);
+ }
+ bid[i * n + _unass_id] = best_i;
+ bid_increments[i * n + _unass_id] = best - better + eps;
+ atomicMax(&max_increments[i * n + best_i], best - better + eps);
+ }
+ }
+}
+
+__global__ void GetMax(int b, int n, int * assignment, int * bid, float * bid_increments, float * max_increments, int * max_idx) {
+ for (int i = blockIdx.x; i < b; i += gridDim.x) {
+ int j = threadIdx.x + blockIdx.y * blockDim.x;
+ if (assignment[i * n + j] == -1) {
+ int bid_id = bid[i * n + j];
+ float bid_inc = bid_increments[i * n + j];
+ float max_inc = max_increments[i * n + bid_id];
+ if (bid_inc - 1e-6 <= max_inc && max_inc <= bid_inc + 1e-6)
+ {
+ max_idx[i * n + bid_id] = j;
+ }
+ }
+ }
+}
+
+__global__ void Assign(int b, int n, int * assignment, int * assignment_inv, float * price, int * bid, float * bid_increments, float * max_increments, int * max_idx, bool last) {
+ for (int i = blockIdx.x; i < b; i += gridDim.x) {
+ int j = threadIdx.x + blockIdx.y * blockDim.x;
+ if (assignment[i * n + j] == -1) {
+ int bid_id = bid[i * n + j];
+ if (last || max_idx[i * n + bid_id] == j)
+ {
+ float bid_inc = bid_increments[i * n + j];
+ int ass_inv = assignment_inv[i * n + bid_id];
+ if (!last && ass_inv != -1) {
+ assignment[i * n + ass_inv] = -1;
+ }
+ assignment_inv[i * n + bid_id] = j;
+ assignment[i * n + j] = bid_id;
+ price[i * n + bid_id] += bid_inc;
+ max_increments[i * n + bid_id] = -1e9;
+ }
+ }
+ }
+}
+
+__global__ void CalcDist(int b, int n, float * xyz1, float * xyz2, float * dist, int * assignment) {
+ for (int i = blockIdx.x; i < b; i += gridDim.x) {
+ int j = threadIdx.x + blockIdx.y * blockDim.x;
+ int k = assignment[i * n + j];
+ float deltax = xyz1[(i * n + j) * 3 + 0] - xyz2[(i * n + k) * 3 + 0];
+ float deltay = xyz1[(i * n + j) * 3 + 1] - xyz2[(i * n + k) * 3 + 1];
+ float deltaz = xyz1[(i * n + j) * 3 + 2] - xyz2[(i * n + k) * 3 + 2];
+ dist[i * n + j] = deltax * deltax + deltay * deltay + deltaz * deltaz;
+ }
+}
+
+int emd_cuda_forward(at::Tensor xyz1, at::Tensor xyz2, at::Tensor dist, at::Tensor assignment, at::Tensor price,
+ at::Tensor assignment_inv, at::Tensor bid, at::Tensor bid_increments, at::Tensor max_increments,
+ at::Tensor unass_idx, at::Tensor unass_cnt, at::Tensor unass_cnt_sum, at::Tensor cnt_tmp, at::Tensor max_idx, float eps, int iters) {
+
+ const auto batch_size = xyz1.size(0);
+ const auto n = xyz1.size(1); //num_points point cloud A
+ const auto m = xyz2.size(1); //num_points point cloud B
+
+ if (n != m) {
+ printf("Input Error! The two point clouds should have the same size.\n");
+ return -1;
+ }
+
+ if (batch_size > 512) {
+ printf("Input Error! The batch size should be less than 512.\n");
+ return -1;
+ }
+
+ if (n % 1024 != 0) {
+ printf("Input Error! The size of the point clouds should be a multiple of 1024.\n");
+ return -1;
+ }
+
+ //cudaEvent_t start,stop;
+ //cudaEventCreate(&start);
+ //cudaEventCreate(&stop);
+ //cudaEventRecord(start);
+ //int iters = 50;
+ for (int i = 0; i < iters; i++) {
+ clear<<<1, batch_size>>>(batch_size, cnt_tmp.data<int>(), unass_cnt.data<int>());
+ calc_unass_cnt<<<dim3(batch_size, n / 1024, 1), 1024>>>(batch_size, n, assignment.data<int>(), unass_cnt.data<int>());
+ calc_unass_cnt_sum<<<1, batch_size>>>(batch_size, unass_cnt.data<int>(), unass_cnt_sum.data<int>());
+ calc_unass_idx<<<dim3(batch_size, n / 1024, 1), 1024>>>(batch_size, n, assignment.data<int>(), unass_idx.data<int>(), unass_cnt.data<int>(),
+ unass_cnt_sum.data<int>(), cnt_tmp.data<int>());
+ Bid<<<dim3(batch_size, n / 1024, 1), 1024>>>(batch_size, n, xyz1.data<float>(), xyz2.data<float>(), eps, assignment.data<int>(), assignment_inv.data<int>(),
+ price.data<float>(), bid.data<int>(), bid_increments.data<float>(), max_increments.data<float>(),
+ unass_cnt.data<int>(), unass_cnt_sum.data<int>(), unass_idx.data<int>());
+ GetMax<<<dim3(batch_size, n / 1024, 1), 1024>>>(batch_size, n, assignment.data<int>(), bid.data<int>(), bid_increments.data<float>(), max_increments.data<float>(), max_idx.data<int>());
+ Assign<<<dim3(batch_size, n / 1024, 1), 1024>>>(batch_size, n, assignment.data<int>(), assignment_inv.data<int>(), price.data<float>(), bid.data<int>(),
+ bid_increments.data<float>(), max_increments.data<float>(), max_idx.data<int>(), i == iters - 1);
+ }
+ CalcDist<<<dim3(batch_size, n / 1024, 1), 1024>>>(batch_size, n, xyz1.data<float>(), xyz2.data<float>(), dist.data<float>(), assignment.data<int>());
+ //cudaEventRecord(stop);
+ //cudaEventSynchronize(stop);
+ //float elapsedTime;
+ //cudaEventElapsedTime(&elapsedTime,start,stop);
+ //printf("%lf\n", elapsedTime);
+
+ cudaError_t err = cudaGetLastError();
+ if (err != cudaSuccess) {
+ printf("error in nnd Output: %s\n", cudaGetErrorString(err));
+ return 0;
+ }
+ return 1;
+}
+
+__global__ void NmDistanceGradKernel(int b, int n, const float * xyz1, const float * xyz2, const float * grad_dist, const int * idx, float * grad_xyz){
+ for (int i = blockIdx.x; i < b; i += gridDim.x) {
+ for (int j = threadIdx.x + blockIdx.y * blockDim.x; j < n; j += blockDim.x * gridDim.y) {
+ float x1 = xyz1[(i * n + j) * 3 + 0];
+ float y1 = xyz1[(i * n + j) * 3 + 1];
+ float z1 = xyz1[(i * n + j) * 3 + 2];
+ int j2 = idx[i * n + j];
+ float x2 = xyz2[(i * n + j2) * 3 + 0];
+ float y2 = xyz2[(i * n + j2) * 3 + 1];
+ float z2 = xyz2[(i * n + j2) * 3 + 2];
+ float g = grad_dist[i * n + j] * 2;
+ atomicAdd(&(grad_xyz[(i * n + j) * 3 + 0]), g * (x1 - x2));
+ atomicAdd(&(grad_xyz[(i * n + j) * 3 + 1]), g * (y1 - y2));
+ atomicAdd(&(grad_xyz[(i * n + j) * 3 + 2]), g * (z1 - z2));
+ }
+ }
+}
+
+int emd_cuda_backward(at::Tensor xyz1, at::Tensor xyz2, at::Tensor gradxyz, at::Tensor graddist, at::Tensor idx){
+ const auto batch_size = xyz1.size(0);
+ const auto n = xyz1.size(1);
+ const auto m = xyz2.size(1);
+
+ NmDistanceGradKernel<<<dim3(batch_size, n / 1024, 1), 1024>>>(batch_size, n, xyz1.data<float>(), xyz2.data<float>(), graddist.data<float>(), idx.data<int>(), gradxyz.data<float>());
+
+ cudaError_t err = cudaGetLastError();
+ if (err != cudaSuccess) {
+ printf("error in nnd get grad: %s\n", cudaGetErrorString(err));
+ return 0;
+ }
+ return 1;
+
+}
diff --git a/lidm/eval/modules/emd/emd_module.py b/lidm/eval/modules/emd/emd_module.py
new file mode 100644
index 0000000000000000000000000000000000000000..3065509eb77942610aab93a85c21f30bf6bf9bed
--- /dev/null
+++ b/lidm/eval/modules/emd/emd_module.py
@@ -0,0 +1,112 @@
+# EMD approximation module (based on auction algorithm)
+# memory complexity: O(n)
+# time complexity: O(n^2 * iter)
+# author: Minghua Liu
+
+# Input:
+# xyz1, xyz2: [#batch, #points, 3]
+# where xyz1 is the predicted point cloud and xyz2 is the ground truth point cloud
+# two point clouds should have same size and be normalized to [0, 1]
+# #points should be a multiple of 1024
+# #batch should be no greater than 512
+# eps is a parameter which balances the error rate and the speed of convergence
+# iters is the number of iteration
+# we only calculate gradient for xyz1
+
+# Output:
+# dist: [#batch, #points], sqrt(dist) -> L2 distance
+# assignment: [#batch, #points], index of the matched point in the ground truth point cloud
+# the result is an approximation and the assignment is not guranteed to be a bijection
+import importlib
+import os
+import time
+import numpy as np
+import torch
+from torch import nn
+from torch.autograd import Function
+
+emd_found = importlib.find_loader("emd") is not None
+if not emd_found:
+ ## Cool trick from https://github.com/chrdiller
+ print("Jitting EMD 3D")
+
+ from torch.utils.cpp_extension import load
+
+ emd = load(name="emd",
+ sources=[
+ "/".join(os.path.abspath(__file__).split('/')[:-1] + ["emd.cpp"]),
+ "/".join(os.path.abspath(__file__).split('/')[:-1] + ["emd_cuda.cu"]),
+ ])
+ print("Loaded JIT 3D CUDA emd")
+else:
+ import emd
+ print("Loaded compiled 3D CUDA emd")
+
+
+class emdFunction(Function):
+ @staticmethod
+ def forward(ctx, xyz1, xyz2, eps, iters):
+ batchsize, n, _ = xyz1.size()
+ _, m, _ = xyz2.size()
+
+ assert (n == m)
+ assert (xyz1.size()[0] == xyz2.size()[0])
+ # assert(n % 1024 == 0)
+ assert (batchsize <= 512)
+
+ xyz1 = xyz1.contiguous().float().cuda()
+ xyz2 = xyz2.contiguous().float().cuda()
+ dist = torch.zeros(batchsize, n, device='cuda').contiguous()
+ assignment = torch.zeros(batchsize, n, device='cuda', dtype=torch.int32).contiguous() - 1
+ assignment_inv = torch.zeros(batchsize, m, device='cuda', dtype=torch.int32).contiguous() - 1
+ price = torch.zeros(batchsize, m, device='cuda').contiguous()
+ bid = torch.zeros(batchsize, n, device='cuda', dtype=torch.int32).contiguous()
+ bid_increments = torch.zeros(batchsize, n, device='cuda').contiguous()
+ max_increments = torch.zeros(batchsize, m, device='cuda').contiguous()
+ unass_idx = torch.zeros(batchsize * n, device='cuda', dtype=torch.int32).contiguous()
+ max_idx = torch.zeros(batchsize * m, device='cuda', dtype=torch.int32).contiguous()
+ unass_cnt = torch.zeros(512, dtype=torch.int32, device='cuda').contiguous()
+ unass_cnt_sum = torch.zeros(512, dtype=torch.int32, device='cuda').contiguous()
+ cnt_tmp = torch.zeros(512, dtype=torch.int32, device='cuda').contiguous()
+
+ emd.forward(xyz1, xyz2, dist, assignment, price, assignment_inv, bid, bid_increments, max_increments, unass_idx,
+ unass_cnt, unass_cnt_sum, cnt_tmp, max_idx, eps, iters)
+
+ ctx.save_for_backward(xyz1, xyz2, assignment)
+ return dist, assignment
+
+ @staticmethod
+ def backward(ctx, graddist, gradidx):
+ xyz1, xyz2, assignment = ctx.saved_tensors
+ graddist = graddist.contiguous()
+
+ gradxyz1 = torch.zeros(xyz1.size(), device='cuda').contiguous()
+ gradxyz2 = torch.zeros(xyz2.size(), device='cuda').contiguous()
+
+ emd.backward(xyz1, xyz2, gradxyz1, graddist, assignment)
+ return gradxyz1, gradxyz2, None, None
+
+
+class emdModule(nn.Module):
+ def __init__(self):
+ super(emdModule, self).__init__()
+
+ def forward(self, input1, input2, eps, iters):
+ return emdFunction.apply(input1, input2, eps, iters)
+
+
+def test_emd():
+ x1 = torch.rand(20, 8192, 3).cuda()
+ x2 = torch.rand(20, 8192, 3).cuda()
+ emd = emdModule()
+ start_time = time.perf_counter()
+ dis, assigment = emd(x1, x2, 0.05, 3000)
+ print("Input_size: ", x1.shape)
+ print("Runtime: %lfs" % (time.perf_counter() - start_time))
+ print("EMD: %lf" % np.sqrt(dis.cpu()).mean())
+ print("|set(assignment)|: %d" % assigment.unique().numel())
+ assigment = assigment.cpu().numpy()
+ assigment = np.expand_dims(assigment, -1)
+ x2 = np.take_along_axis(x2, assigment, axis=1)
+ d = (x1 - x2) * (x1 - x2)
+ print("Verified EMD: %lf" % np.sqrt(d.cpu().sum(-1)).mean())
diff --git a/lidm/eval/modules/emd/setup.py b/lidm/eval/modules/emd/setup.py
new file mode 100644
index 0000000000000000000000000000000000000000..8588de957bc285372be5a59a6e086af4954dc99b
--- /dev/null
+++ b/lidm/eval/modules/emd/setup.py
@@ -0,0 +1,14 @@
+from setuptools import setup
+from torch.utils.cpp_extension import BuildExtension, CUDAExtension
+
+setup(
+ name='emd',
+ ext_modules=[
+ CUDAExtension('emd', [
+ 'emd.cpp',
+ 'emd_cuda.cu',
+ ]),
+ ],
+ cmdclass={
+ 'build_ext': BuildExtension
+ })
\ No newline at end of file
diff --git a/lidm/models/__init__.py b/lidm/models/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/lidm/models/autoencoder.py b/lidm/models/autoencoder.py
new file mode 100644
index 0000000000000000000000000000000000000000..089130695164fbb30460439c78efd96abcae33e5
--- /dev/null
+++ b/lidm/models/autoencoder.py
@@ -0,0 +1,465 @@
+import numpy as np
+import torch
+import pytorch_lightning as pl
+import torch.nn.functional as F
+from contextlib import contextmanager
+
+from taming.modules.vqvae.quantize import VectorQuantizer2 as VectorQuantizer
+
+from ..modules.diffusion import model_lidm, model_ldm
+from ..modules.distributions.distributions import DiagonalGaussianDistribution
+from ..modules.ema import LitEma
+from ..utils.misc_utils import instantiate_from_config
+
+
+class VQModel(pl.LightningModule):
+ def __init__(self,
+ ddconfig,
+ n_embed,
+ embed_dim,
+ lossconfig=None,
+ ckpt_path=None,
+ ignore_keys=[],
+ image_key="image",
+ colorize_nlabels=None,
+ monitor=None,
+ batch_resize_range=None,
+ scheduler_config=None,
+ lr_g_factor=1.0,
+ remap=None,
+ sane_index_shape=False, # tell vector quantizer to return indices as bhw
+ use_ema=False,
+ lib_name='ldm',
+ use_mask=False,
+ **kwargs
+ ):
+ super().__init__()
+ self.embed_dim = embed_dim
+ self.n_embed = n_embed
+ self.image_key = image_key
+ self.use_mask = use_mask
+ model_lib = eval(f'model_{lib_name}')
+ self.encoder = model_lib.Encoder(**ddconfig)
+ self.decoder = model_lib.Decoder(**ddconfig)
+ if lossconfig is not None:
+ self.loss = instantiate_from_config(lossconfig)
+ self.quantize = VectorQuantizer(n_embed, embed_dim, beta=0.25,
+ remap=remap,
+ sane_index_shape=sane_index_shape)
+ self.quant_conv = torch.nn.Conv2d(ddconfig["z_channels"], embed_dim, 1)
+ self.post_quant_conv = torch.nn.Conv2d(embed_dim, ddconfig["z_channels"], 1)
+ if colorize_nlabels is not None:
+ assert type(colorize_nlabels) == int
+ self.register_buffer("colorize", torch.randn(3, colorize_nlabels, 1, 1))
+ if monitor is not None:
+ self.monitor = monitor
+ self.batch_resize_range = batch_resize_range
+ if self.batch_resize_range is not None:
+ print(f"{self.__class__.__name__}: Using per-batch resizing in range {batch_resize_range}.")
+
+ self.use_ema = use_ema
+ if self.use_ema:
+ self.model_ema = LitEma(self)
+ print(f"Keeping EMAs of {len(list(self.model_ema.buffers()))}.")
+
+ if ckpt_path is not None:
+ self.init_from_ckpt(ckpt_path, ignore_keys=ignore_keys)
+ self.scheduler_config = scheduler_config
+ self.lr_g_factor = lr_g_factor
+
+ @contextmanager
+ def ema_scope(self, context=None):
+ if self.use_ema:
+ self.model_ema.store(self.parameters())
+ self.model_ema.copy_to(self)
+ if context is not None:
+ print(f"{context}: Switched to EMA weights")
+ try:
+ yield None
+ finally:
+ if self.use_ema:
+ self.model_ema.restore(self.parameters())
+ if context is not None:
+ print(f"{context}: Restored training weights")
+
+ def init_from_ckpt(self, path, ignore_keys=list()):
+ sd = torch.load(path, map_location="cpu")["state_dict"]
+ keys = list(sd.keys())
+ for k in keys:
+ for ik in ignore_keys:
+ if k.startswith(ik):
+ print("Deleting key {} from state_dict.".format(k))
+ del sd[k]
+ missing, unexpected = self.load_state_dict(sd, strict=False)
+ print(f"Restored from {path} with {len(missing)} missing and {len(unexpected)} unexpected keys")
+ if len(missing) > 0:
+ print(f"Missing Keys: {missing}")
+ print(f"Unexpected Keys: {unexpected}")
+
+ def on_train_batch_end(self, *args, **kwargs):
+ if self.use_ema:
+ self.model_ema(self)
+
+ def encode(self, x):
+ h = self.encoder(x)
+ h = self.quant_conv(h)
+ quant, emb_loss, info = self.quantize(h)
+ return quant, emb_loss, info
+
+ def encode_to_prequant(self, x):
+ h = self.encoder(x)
+ h = self.quant_conv(h)
+ return h
+
+ def decode(self, quant):
+ quant = self.post_quant_conv(quant)
+ dec = self.decoder(quant)
+ return dec
+
+ def decode_code(self, code_b):
+ quant_b = self.quantize.embed_code(code_b)
+ dec = self.decode(quant_b)
+ return dec
+
+ def forward(self, input, return_pred_indices=False):
+ quant, diff, (_, _, ind) = self.encode(input)
+ dec = self.decode(quant)
+ if return_pred_indices:
+ return dec, diff, ind
+ return dec, diff
+
+ def get_input(self, batch, k):
+ x = batch[k]
+ # if len(x.shape) == 3:
+ # x = x[..., None]
+
+ if self.batch_resize_range is not None:
+ lower_size = self.batch_resize_range[0]
+ upper_size = self.batch_resize_range[1]
+ if self.global_step <= 4:
+ # do the first few batches with max size to avoid later oom
+ new_resize = upper_size
+ else:
+ new_resize = np.random.choice(np.arange(lower_size, upper_size + 16, 16))
+ if new_resize != x.shape[2]:
+ x = F.interpolate(x, size=new_resize, mode="bicubic")
+ x = x.detach()
+ return x
+
+ def get_mask(self, batch):
+ mask = batch['mask']
+ # if len(mask.shape) == 3:
+ # mask = mask[..., None]
+ return mask
+
+ def training_step(self, batch, batch_idx, optimizer_idx):
+ # https://github.com/pytorch/pytorch/issues/37142
+ # try not to fool the heuristics
+ x = self.get_input(batch, self.image_key)
+ m = self.get_mask(batch) if self.use_mask else None
+ x_rec, qloss, ind = self(x, return_pred_indices=True)
+
+ if optimizer_idx == 0:
+ # autoencoder
+ aeloss, log_dict_ae = self.loss(qloss, x, x_rec, optimizer_idx, self.global_step,
+ last_layer=self.get_last_layer(), split="train",
+ predicted_indices=None, masks=m)
+ self.log_dict(log_dict_ae, prog_bar=False, logger=True, on_step=True, on_epoch=True)
+ return aeloss
+
+ if optimizer_idx == 1:
+ # discriminator
+ discloss, log_dict_disc = self.loss(qloss, x, x_rec, optimizer_idx, self.global_step,
+ last_layer=self.get_last_layer(), split="train",
+ masks=m)
+ self.log_dict(log_dict_disc, prog_bar=False, logger=True, on_step=True, on_epoch=True)
+ return discloss
+
+ def validation_step(self, batch, batch_idx):
+ log_dict = self._validation_step(batch, batch_idx)
+ if self.use_ema:
+ with self.ema_scope():
+ log_dict_ema = self._validation_step(batch, batch_idx, suffix="_ema")
+ return log_dict
+
+ def _validation_step(self, batch, batch_idx, suffix=""):
+ x = self.get_input(batch, self.image_key)
+ m = self.get_mask(batch) if self.use_mask else None
+ xrec, qloss, ind = self(x, return_pred_indices=True)
+ aeloss, log_dict_ae = self.loss(qloss, x, xrec, 0,
+ self.global_step,
+ last_layer=self.get_last_layer(),
+ split="val" + suffix,
+ predicted_indices=None,
+ masks=m
+ )
+
+ discloss, log_dict_disc = self.loss(qloss, x, xrec, 1,
+ self.global_step,
+ last_layer=self.get_last_layer(),
+ split="val" + suffix,
+ predicted_indices=None,
+ masks=m
+ )
+ rec_loss = log_dict_ae[f"val{suffix}/rec_loss"]
+ self.log(f"val{suffix}/rec_loss", rec_loss,
+ prog_bar=True, logger=True, on_step=False, on_epoch=True, sync_dist=True)
+ self.log(f"val{suffix}/aeloss", aeloss,
+ prog_bar=True, logger=True, on_step=False, on_epoch=True, sync_dist=True)
+ del log_dict_ae[f"val{suffix}/rec_loss"]
+ self.log_dict(log_dict_ae)
+ self.log_dict(log_dict_disc)
+ return self.log_dict
+
+ def configure_optimizers(self):
+ lr_d = self.learning_rate
+ lr_g = self.lr_g_factor * self.learning_rate
+ # print("lr_d", lr_d)
+ # print("lr_g", lr_g)
+ opt_ae = torch.optim.Adam(list(self.encoder.parameters()) +
+ list(self.decoder.parameters()) +
+ list(self.quantize.parameters()) +
+ list(self.quant_conv.parameters()) +
+ list(self.post_quant_conv.parameters()),
+ lr=lr_g, betas=(0.5, 0.9))
+ opt_disc = torch.optim.Adam(self.loss.discriminator.parameters(),
+ lr=lr_d, betas=(0.5, 0.9))
+
+ if self.scheduler_config is not None:
+ scheduler = instantiate_from_config(self.scheduler_config)
+
+ print("Setting up LambdaLR scheduler...")
+ scheduler = [
+ {
+ 'scheduler': LambdaLR(opt_ae, lr_lambda=scheduler.schedule),
+ 'interval': 'step',
+ 'frequency': 1
+ },
+ {
+ 'scheduler': LambdaLR(opt_disc, lr_lambda=scheduler.schedule),
+ 'interval': 'step',
+ 'frequency': 1
+ },
+ ]
+ return [opt_ae, opt_disc], scheduler
+ return [opt_ae, opt_disc], []
+
+ def get_last_layer(self):
+ return self.decoder.conv_out.weight
+
+ @torch.no_grad()
+ def log_images(self, batch, only_inputs=False, plot_ema=False, **kwargs):
+ log = dict()
+ x = self.get_input(batch, self.image_key)
+ x = x.to(self.device)
+ if only_inputs:
+ log["inputs"] = x
+ return log
+ xrec, _ = self(x)
+ if self.use_mask:
+ mask = xrec[:, 1:2] < 0.
+ xrec = xrec[:, 0:1]
+ xrec[mask] = -1.
+ log["inputs"] = x
+ log["reconstructions"] = xrec
+ if plot_ema:
+ with self.ema_scope():
+ xrec_ema, _ = self(x)
+ log["reconstructions_ema"] = xrec_ema
+ return log
+
+ def to_rgb(self, x):
+ assert self.image_key == "segmentation"
+ if not hasattr(self, "colorize"):
+ self.register_buffer("colorize", torch.randn(3, x.shape[1], 1, 1).to(x))
+ x = F.conv2d(x, weight=self.colorize)
+ x = 2. * (x - x.min()) / (x.max() - x.min()) - 1.
+ return x
+
+
+class VQModelInterface(VQModel):
+ def __init__(self, embed_dim, *args, **kwargs):
+ super().__init__(embed_dim=embed_dim, *args, **kwargs)
+ self.embed_dim = embed_dim
+
+ def encode(self, x):
+ h = self.encoder(x)
+ h = self.quant_conv(h)
+ return h
+
+ def decode(self, h, force_not_quantize=False):
+ # also go through quantization layer
+ if not force_not_quantize:
+ quant, emb_loss, info = self.quantize(h)
+ else:
+ quant = h
+ quant = self.post_quant_conv(quant)
+ dec = self.decoder(quant)
+ if self.use_mask:
+ mask = dec[:, 1:2] < 0.
+ dec = dec[:, 0:1]
+ dec[mask] = -1.
+ return dec
+
+
+class AutoencoderKL(pl.LightningModule):
+ def __init__(self,
+ ddconfig,
+ lossconfig,
+ embed_dim,
+ ckpt_path=None,
+ ignore_keys=[],
+ image_key="image",
+ colorize_nlabels=None,
+ monitor=None,
+ lib_name='ldm',
+ use_mask=False
+ ):
+ super().__init__()
+ self.image_key = image_key
+ self.use_mask = use_mask
+ model_lib = eval(f'model_{lib_name}')
+ self.encoder = model_lib.Encoder(**ddconfig)
+ self.decoder = model_lib.Decoder(**ddconfig)
+ self.loss = instantiate_from_config(lossconfig)
+ assert ddconfig["double_z"]
+ self.quant_conv = torch.nn.Conv2d(2 * ddconfig["z_channels"], 2 * embed_dim, 1)
+ self.post_quant_conv = torch.nn.Conv2d(embed_dim, ddconfig["z_channels"], 1)
+ self.embed_dim = embed_dim
+ if colorize_nlabels is not None:
+ assert type(colorize_nlabels) == int
+ self.register_buffer("colorize", torch.randn(3, colorize_nlabels, 1, 1))
+ if monitor is not None:
+ self.monitor = monitor
+ if ckpt_path is not None:
+ self.init_from_ckpt(ckpt_path, ignore_keys=ignore_keys)
+
+ def init_from_ckpt(self, path, ignore_keys=list()):
+ sd = torch.load(path, map_location="cpu")["state_dict"]
+ keys = list(sd.keys())
+ for k in keys:
+ for ik in ignore_keys:
+ if k.startswith(ik):
+ print("Deleting key {} from state_dict.".format(k))
+ del sd[k]
+ self.load_state_dict(sd, strict=False)
+ print(f"Restored from {path}")
+
+ def encode(self, x):
+ h = self.encoder(x)
+ moments = self.quant_conv(h)
+ posterior = DiagonalGaussianDistribution(moments)
+ return posterior
+
+ def decode(self, z):
+ z = self.post_quant_conv(z)
+ dec = self.decoder(z)
+ return dec
+
+ def forward(self, input, sample_posterior=True):
+ posterior = self.encode(input)
+ if sample_posterior:
+ z = posterior.sample()
+ else:
+ z = posterior.mode()
+ dec = self.decode(z)
+ return dec, posterior
+
+ def get_input(self, batch, k):
+ x = batch[k]
+ if len(x.shape) == 3:
+ x = x[:, None]
+ return x
+
+ def training_step(self, batch, batch_idx, optimizer_idx):
+ inputs = self.get_input(batch, self.image_key)
+ reconstructions, posterior = self(inputs)
+
+ if optimizer_idx == 0:
+ # train encoder+decoder+logvar
+ aeloss, log_dict_ae = self.loss(inputs, reconstructions, posterior, optimizer_idx, self.global_step,
+ last_layer=self.get_last_layer(), split="train")
+ self.log("aeloss", aeloss, prog_bar=True, logger=True, on_step=True, on_epoch=True)
+ self.log_dict(log_dict_ae, prog_bar=False, logger=True, on_step=True, on_epoch=False)
+ return aeloss
+
+ if optimizer_idx == 1:
+ # train the discriminator
+ discloss, log_dict_disc = self.loss(inputs, reconstructions, posterior, optimizer_idx, self.global_step,
+ last_layer=self.get_last_layer(), split="train")
+
+ self.log("discloss", discloss, prog_bar=True, logger=True, on_step=True, on_epoch=True)
+ self.log_dict(log_dict_disc, prog_bar=False, logger=True, on_step=True, on_epoch=False)
+ return discloss
+
+ def validation_step(self, batch, batch_idx):
+ inputs = self.get_input(batch, self.image_key)
+ reconstructions, posterior = self(inputs)
+ aeloss, log_dict_ae = self.loss(inputs, reconstructions, posterior, 0, self.global_step,
+ last_layer=self.get_last_layer(), split="val")
+ discloss, log_dict_disc = self.loss(inputs, reconstructions, posterior, 1, self.global_step,
+ last_layer=self.get_last_layer(), split="val")
+
+ self.log("val/rec_loss", log_dict_ae["val/rec_loss"])
+ self.log_dict(log_dict_ae)
+ self.log_dict(log_dict_disc)
+ return self.log_dict
+
+ def configure_optimizers(self):
+ lr = self.learning_rate
+ opt_ae = torch.optim.Adam(list(self.encoder.parameters()) +
+ list(self.decoder.parameters()) +
+ list(self.quant_conv.parameters()) +
+ list(self.post_quant_conv.parameters()),
+ lr=lr, betas=(0.5, 0.9))
+ opt_disc = torch.optim.Adam(self.loss.discriminator.parameters(),
+ lr=lr, betas=(0.5, 0.9))
+ return [opt_ae, opt_disc], []
+
+ def get_last_layer(self):
+ return self.decoder.conv_out.weight
+
+ @torch.no_grad()
+ def log_images(self, batch, only_inputs=False, **kwargs):
+ log = dict()
+ x = self.get_input(batch, self.image_key)
+ x = x.to(self.device)
+ if not only_inputs:
+ xrec, posterior = self(x)
+ if x.shape[1] > 3:
+ # colorize with random projection
+ assert xrec.shape[1] > 3
+ x = self.to_rgb(x)
+ xrec = self.to_rgb(xrec)
+ log["samples"] = self.decode(torch.randn_like(posterior.sample()))
+ log["reconstructions"] = xrec
+ log["inputs"] = x
+ return log
+
+ def to_rgb(self, x):
+ assert self.image_key == "segmentation"
+ if not hasattr(self, "colorize"):
+ self.register_buffer("colorize", torch.randn(3, x.shape[1], 1, 1).to(x))
+ x = F.conv2d(x, weight=self.colorize)
+ x = 2. * (x - x.min()) / (x.max() - x.min()) - 1.
+ return x
+
+
+class IdentityFirstStage(torch.nn.Module):
+ def __init__(self, *args, vq_interface=False, **kwargs):
+ self.vq_interface = vq_interface
+ super().__init__()
+
+ def encode(self, x, *args, **kwargs):
+ return x
+
+ def decode(self, x, *args, **kwargs):
+ return x
+
+ def quantize(self, x, *args, **kwargs):
+ if self.vq_interface:
+ return x, None, [None, None, None]
+ return x
+
+ def forward(self, x, *args, **kwargs):
+ return x
diff --git a/lidm/models/diffusion/__init__.py b/lidm/models/diffusion/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/lidm/models/diffusion/classifier.py b/lidm/models/diffusion/classifier.py
new file mode 100644
index 0000000000000000000000000000000000000000..82d2d6d0ef132a83e25bc84e324252a9aa1d513b
--- /dev/null
+++ b/lidm/models/diffusion/classifier.py
@@ -0,0 +1,267 @@
+import os
+import torch
+import pytorch_lightning as pl
+from omegaconf import OmegaConf
+from torch.nn import functional as F
+from torch.optim import AdamW
+from torch.optim.lr_scheduler import LambdaLR
+from copy import deepcopy
+from einops import rearrange
+from glob import glob
+from natsort import natsorted
+
+from ...modules.diffusion.openaimodel import EncoderUNetModel, UNetModel
+from ...utils.misc_utils import log_txt_as_img, default, ismap, instantiate_from_config
+
+__models__ = {
+ 'class_label': EncoderUNetModel,
+ 'segmentation': UNetModel
+}
+
+
+def disabled_train(self, mode=True):
+ """Overwrite model.train with this function to make sure train/eval mode
+ does not change anymore."""
+ return self
+
+
+class NoisyLatentImageClassifier(pl.LightningModule):
+
+ def __init__(self,
+ diffusion_path,
+ num_classes,
+ ckpt_path=None,
+ pool='attention',
+ label_key=None,
+ diffusion_ckpt_path=None,
+ scheduler_config=None,
+ weight_decay=1.e-2,
+ log_steps=10,
+ monitor='val/loss',
+ *args,
+ **kwargs):
+ super().__init__(*args, **kwargs)
+ self.num_classes = num_classes
+ # get latest config of diffusion model
+ diffusion_config = natsorted(glob(os.path.join(diffusion_path, 'configs', '*-project.yaml')))[-1]
+ self.diffusion_config = OmegaConf.load(diffusion_config).model
+ self.diffusion_config.params.ckpt_path = diffusion_ckpt_path
+ self.load_diffusion()
+
+ self.monitor = monitor
+ self.numd = self.diffusion_model.first_stage_model.encoder.num_resolutions - 1
+ self.log_time_interval = self.diffusion_model.num_timesteps // log_steps
+ self.log_steps = log_steps
+
+ self.label_key = label_key if not hasattr(self.diffusion_model, 'cond_stage_key') \
+ else self.diffusion_model.cond_stage_key
+
+ assert self.label_key is not None, 'label_key neither in diffusion model nor in model.params'
+
+ if self.label_key not in __models__:
+ raise NotImplementedError()
+
+ self.load_classifier(ckpt_path, pool)
+
+ self.scheduler_config = scheduler_config
+ self.use_scheduler = self.scheduler_config is not None
+ self.weight_decay = weight_decay
+
+ def init_from_ckpt(self, path, ignore_keys=list(), only_model=False):
+ sd = torch.load(path, map_location="cpu")
+ if "state_dict" in list(sd.keys()):
+ sd = sd["state_dict"]
+ keys = list(sd.keys())
+ for k in keys:
+ for ik in ignore_keys:
+ if k.startswith(ik):
+ print("Deleting key {} from state_dict.".format(k))
+ del sd[k]
+ missing, unexpected = self.load_state_dict(sd, strict=False) if not only_model else self.model.load_state_dict(
+ sd, strict=False)
+ print(f"Restored from {path} with {len(missing)} missing and {len(unexpected)} unexpected keys")
+ if len(missing) > 0:
+ print(f"Missing Keys: {missing}")
+ if len(unexpected) > 0:
+ print(f"Unexpected Keys: {unexpected}")
+
+ def load_diffusion(self):
+ model = instantiate_from_config(self.diffusion_config)
+ self.diffusion_model = model.eval()
+ self.diffusion_model.train = disabled_train
+ for param in self.diffusion_model.parameters():
+ param.requires_grad = False
+
+ def load_classifier(self, ckpt_path, pool):
+ model_config = deepcopy(self.diffusion_config.params.unet_config.params)
+ model_config.in_channels = self.diffusion_config.params.unet_config.params.out_channels
+ model_config.out_channels = self.num_classes
+ if self.label_key == 'class_label':
+ model_config.pool = pool
+
+ self.model = __models__[self.label_key](**model_config)
+ if ckpt_path is not None:
+ print('#####################################################################')
+ print(f'load from ckpt "{ckpt_path}"')
+ print('#####################################################################')
+ self.init_from_ckpt(ckpt_path)
+
+ @torch.no_grad()
+ def get_x_noisy(self, x, t, noise=None):
+ noise = default(noise, lambda: torch.randn_like(x))
+ continuous_sqrt_alpha_cumprod = None
+ if self.diffusion_model.use_continuous_noise:
+ continuous_sqrt_alpha_cumprod = self.diffusion_model.sample_continuous_noise_level(x.shape[0], t + 1)
+ # todo: make sure t+1 is correct here
+
+ return self.diffusion_model.q_sample(x_start=x, t=t, noise=noise,
+ continuous_sqrt_alpha_cumprod=continuous_sqrt_alpha_cumprod)
+
+ def forward(self, x_noisy, t, *args, **kwargs):
+ return self.model(x_noisy, t)
+
+ @torch.no_grad()
+ def get_input(self, batch, k):
+ x = batch[k]
+ if len(x.shape) == 3:
+ x = x[..., None]
+ x = rearrange(x, 'b h w c -> b c h w')
+ x = x.to(memory_format=torch.contiguous_format).float()
+ return x
+
+ @torch.no_grad()
+ def get_conditioning(self, batch, k=None):
+ if k is None:
+ k = self.label_key
+ assert k is not None, 'Needs to provide label key'
+
+ targets = batch[k].to(self.device)
+
+ if self.label_key == 'segmentation':
+ targets = rearrange(targets, 'b h w c -> b c h w')
+ for down in range(self.numd):
+ h, w = targets.shape[-2:]
+ targets = F.interpolate(targets, size=(h // 2, w // 2), mode='nearest')
+
+ # targets = rearrange(targets,'b c h w -> b h w c')
+
+ return targets
+
+ def compute_top_k(self, logits, labels, k, reduction="mean"):
+ _, top_ks = torch.topk(logits, k, dim=1)
+ if reduction == "mean":
+ return (top_ks == labels[:, None]).float().sum(dim=-1).mean().item()
+ elif reduction == "none":
+ return (top_ks == labels[:, None]).float().sum(dim=-1)
+
+ def on_train_epoch_start(self):
+ # save some memory
+ self.diffusion_model.model.to('cpu')
+
+ @torch.no_grad()
+ def write_logs(self, loss, logits, targets):
+ log_prefix = 'train' if self.training else 'val'
+ log = {}
+ log[f"{log_prefix}/loss"] = loss.mean()
+ log[f"{log_prefix}/acc@1"] = self.compute_top_k(
+ logits, targets, k=1, reduction="mean"
+ )
+ log[f"{log_prefix}/acc@5"] = self.compute_top_k(
+ logits, targets, k=5, reduction="mean"
+ )
+
+ self.log_dict(log, prog_bar=False, logger=True, on_step=self.training, on_epoch=True)
+ self.log('loss', log[f"{log_prefix}/loss"], prog_bar=True, logger=False)
+ self.log('global_step', self.global_step, logger=False, on_epoch=False, prog_bar=True)
+ lr = self.optimizers().param_groups[0]['lr']
+ self.log('lr_abs', lr, on_step=True, logger=True, on_epoch=False, prog_bar=True)
+
+ def shared_step(self, batch, t=None):
+ x, *_ = self.diffusion_model.get_input(batch, k=self.diffusion_model.first_stage_key)
+ targets = self.get_conditioning(batch)
+ if targets.dim() == 4:
+ targets = targets.argmax(dim=1)
+ if t is None:
+ t = torch.randint(0, self.diffusion_model.num_timesteps, (x.shape[0],), device=self.device).long()
+ else:
+ t = torch.full(size=(x.shape[0],), fill_value=t, device=self.device).long()
+ x_noisy = self.get_x_noisy(x, t)
+ logits = self(x_noisy, t)
+
+ loss = F.cross_entropy(logits, targets, reduction='none')
+
+ self.write_logs(loss.detach(), logits.detach(), targets.detach())
+
+ loss = loss.mean()
+ return loss, logits, x_noisy, targets
+
+ def training_step(self, batch, batch_idx):
+ loss, *_ = self.shared_step(batch)
+ return loss
+
+ def reset_noise_accs(self):
+ self.noisy_acc = {t: {'acc@1': [], 'acc@5': []} for t in
+ range(0, self.diffusion_model.num_timesteps, self.diffusion_model.log_every_t)}
+
+ def on_validation_start(self):
+ self.reset_noise_accs()
+
+ @torch.no_grad()
+ def validation_step(self, batch, batch_idx):
+ loss, *_ = self.shared_step(batch)
+
+ for t in self.noisy_acc:
+ _, logits, _, targets = self.shared_step(batch, t)
+ self.noisy_acc[t]['acc@1'].append(self.compute_top_k(logits, targets, k=1, reduction='mean'))
+ self.noisy_acc[t]['acc@5'].append(self.compute_top_k(logits, targets, k=5, reduction='mean'))
+
+ return loss
+
+ def configure_optimizers(self):
+ optimizer = AdamW(self.model.parameters(), lr=self.learning_rate, weight_decay=self.weight_decay)
+
+ if self.use_scheduler:
+ scheduler = instantiate_from_config(self.scheduler_config)
+
+ print("Setting up LambdaLR scheduler...")
+ scheduler = [
+ {
+ 'scheduler': LambdaLR(optimizer, lr_lambda=scheduler.schedule),
+ 'interval': 'step',
+ 'frequency': 1
+ }]
+ return [optimizer], scheduler
+
+ return optimizer
+
+ @torch.no_grad()
+ def log_images(self, batch, N=8, *args, **kwargs):
+ log = dict()
+ x = self.get_input(batch, self.diffusion_model.first_stage_key)
+ log['inputs'] = x
+
+ y = self.get_conditioning(batch)
+
+ if self.label_key == 'class_label':
+ y = log_txt_as_img((x.shape[2], x.shape[3]), batch["human_label"])
+ log['labels'] = y
+
+ if ismap(y):
+ log['labels'] = self.diffusion_model.to_rgb(y)
+
+ for step in range(self.log_steps):
+ current_time = step * self.log_time_interval
+
+ _, logits, x_noisy, _ = self.shared_step(batch, t=current_time)
+
+ log[f'inputs@t{current_time}'] = x_noisy
+
+ pred = F.one_hot(logits.argmax(dim=1), num_classes=self.num_classes)
+ pred = rearrange(pred, 'b h w c -> b c h w')
+
+ log[f'pred@t{current_time}'] = self.diffusion_model.to_rgb(pred)
+
+ for key in log:
+ log[key] = log[key][:N]
+
+ return log
diff --git a/lidm/models/diffusion/ddim.py b/lidm/models/diffusion/ddim.py
new file mode 100644
index 0000000000000000000000000000000000000000..8c49dd42cac10dd1743eab581557633124e95826
--- /dev/null
+++ b/lidm/models/diffusion/ddim.py
@@ -0,0 +1,204 @@
+"""SAMPLING ONLY."""
+
+import torch
+import numpy as np
+from tqdm import tqdm
+
+from ...modules.basic import make_ddim_sampling_parameters, make_ddim_timesteps, noise_like
+from ...utils.misc_utils import print_fn
+
+
+class DDIMSampler(object):
+ def __init__(self, model, schedule="linear", **kwargs):
+ super().__init__()
+ self.model = model
+ self.ddpm_num_timesteps = model.num_timesteps
+ self.schedule = schedule
+
+ def register_buffer(self, name, attr):
+ if type(attr) == torch.Tensor:
+ if attr.device != torch.device("cuda"):
+ attr = attr.to(torch.device("cuda"))
+ setattr(self, name, attr)
+
+ def make_schedule(self, ddim_num_steps, ddim_discretize="uniform", ddim_eta=0., verbose=False):
+ self.ddim_timesteps = make_ddim_timesteps(ddim_discr_method=ddim_discretize, num_ddim_timesteps=ddim_num_steps,
+ num_ddpm_timesteps=self.ddpm_num_timesteps, verbose=verbose)
+ alphas_cumprod = self.model.alphas_cumprod
+ assert alphas_cumprod.shape[0] == self.ddpm_num_timesteps, 'alphas have to be defined for each timestep'
+ to_torch = lambda x: x.clone().detach().to(torch.float32).to(self.model.device)
+
+ self.register_buffer('betas', to_torch(self.model.betas))
+ self.register_buffer('alphas_cumprod', to_torch(alphas_cumprod))
+ self.register_buffer('alphas_cumprod_prev', to_torch(self.model.alphas_cumprod_prev))
+
+ # calculations for diffusion q(x_t | x_{t-1}) and others
+ self.register_buffer('sqrt_alphas_cumprod', to_torch(np.sqrt(alphas_cumprod.cpu())))
+ self.register_buffer('sqrt_one_minus_alphas_cumprod', to_torch(np.sqrt(1. - alphas_cumprod.cpu())))
+ self.register_buffer('log_one_minus_alphas_cumprod', to_torch(np.log(1. - alphas_cumprod.cpu())))
+ self.register_buffer('sqrt_recip_alphas_cumprod', to_torch(np.sqrt(1. / alphas_cumprod.cpu())))
+ self.register_buffer('sqrt_recipm1_alphas_cumprod', to_torch(np.sqrt(1. / alphas_cumprod.cpu() - 1)))
+
+ # ddim sampling parameters
+ ddim_sigmas, ddim_alphas, ddim_alphas_prev = make_ddim_sampling_parameters(alphacums=alphas_cumprod.cpu(),
+ ddim_timesteps=self.ddim_timesteps,
+ eta=ddim_eta, verbose=verbose)
+ self.register_buffer('ddim_sigmas', ddim_sigmas)
+ self.register_buffer('ddim_alphas', ddim_alphas)
+ self.register_buffer('ddim_alphas_prev', ddim_alphas_prev)
+ self.register_buffer('ddim_sqrt_one_minus_alphas', np.sqrt(1. - ddim_alphas))
+ sigmas_for_original_sampling_steps = ddim_eta * torch.sqrt(
+ (1 - self.alphas_cumprod_prev) / (1 - self.alphas_cumprod) * (
+ 1 - self.alphas_cumprod / self.alphas_cumprod_prev))
+ self.register_buffer('ddim_sigmas_for_original_num_steps', sigmas_for_original_sampling_steps)
+
+ @torch.no_grad()
+ def sample(self,
+ S,
+ batch_size,
+ shape,
+ conditioning=None,
+ callback=None,
+ normals_sequence=None,
+ img_callback=None,
+ quantize_x0=False,
+ eta=0.,
+ mask=None,
+ x0=None,
+ temperature=1.,
+ noise_dropout=0.,
+ score_corrector=None,
+ corrector_kwargs=None,
+ verbose=False,
+ disable_tqdm=True,
+ x_T=None,
+ log_every_t=100,
+ unconditional_guidance_scale=1.,
+ unconditional_conditioning=None,
+ # this has to come in the same format as the conditioning, # e.g. as encoded tokens, ...
+ **kwargs
+ ):
+ if conditioning is not None:
+ if isinstance(conditioning, dict):
+ cbs = conditioning[list(conditioning.keys())[0]].shape[0]
+ if cbs != batch_size:
+ print(f"Warning: Got {cbs} conditionings but batch-size is {batch_size}")
+ else:
+ if conditioning.shape[0] != batch_size:
+ print(f"Warning: Got {conditioning.shape[0]} conditionings but batch-size is {batch_size}")
+
+ self.make_schedule(ddim_num_steps=S, ddim_eta=eta, verbose=verbose)
+ # sampling
+ C, H, W = shape
+ size = (batch_size, C, H, W)
+ print_fn(f'Data shape for DDIM sampling is {size}, eta {eta}', verbose)
+
+ samples, intermediates = self.ddim_sampling(conditioning, size,
+ callback=callback,
+ img_callback=img_callback,
+ quantize_denoised=quantize_x0,
+ mask=mask, x0=x0,
+ ddim_use_original_steps=False,
+ noise_dropout=noise_dropout,
+ temperature=temperature,
+ score_corrector=score_corrector,
+ corrector_kwargs=corrector_kwargs,
+ x_T=x_T,
+ log_every_t=log_every_t,
+ unconditional_guidance_scale=unconditional_guidance_scale,
+ unconditional_conditioning=unconditional_conditioning,
+ verbose=verbose, disable_tqdm=disable_tqdm)
+ return samples, intermediates
+
+ @torch.no_grad()
+ def ddim_sampling(self, cond, shape,
+ x_T=None, ddim_use_original_steps=False,
+ callback=None, timesteps=None, quantize_denoised=False,
+ mask=None, x0=None, img_callback=None, log_every_t=100,
+ temperature=1., noise_dropout=0., score_corrector=None, corrector_kwargs=None,
+ unconditional_guidance_scale=1., unconditional_conditioning=None, verbose=False, disable_tqdm=True):
+ device = self.model.betas.device
+ b = shape[0]
+ if x_T is None:
+ img = torch.randn(shape, device=device)
+ else:
+ img = x_T
+
+ if timesteps is None:
+ timesteps = self.ddpm_num_timesteps if ddim_use_original_steps else self.ddim_timesteps
+ elif timesteps is not None and not ddim_use_original_steps:
+ subset_end = int(min(timesteps / self.ddim_timesteps.shape[0], 1) * self.ddim_timesteps.shape[0]) - 1
+ timesteps = self.ddim_timesteps[:subset_end]
+
+ intermediates = {'x_inter': [img], 'pred_x0': [img]}
+ time_range = reversed(range(0, timesteps)) if ddim_use_original_steps else np.flip(timesteps)
+ total_steps = timesteps if ddim_use_original_steps else timesteps.shape[0]
+ print_fn(f"Running DDIM Sampling with {total_steps} timesteps", verbose)
+
+ iterator = tqdm(time_range, desc='DDIM Sampler', total=total_steps, disable=disable_tqdm)
+
+ for i, step in enumerate(iterator):
+ index = total_steps - i - 1
+ ts = torch.full((b,), step, device=device, dtype=torch.long)
+
+ if mask is not None:
+ assert x0 is not None
+ img_orig = self.model.q_sample(x0, ts)
+ img = img_orig * mask + (1. - mask) * img
+
+ outs = self.p_sample_ddim(img, cond, ts, index=index, use_original_steps=ddim_use_original_steps,
+ quantize_denoised=quantize_denoised, temperature=temperature,
+ noise_dropout=noise_dropout, score_corrector=score_corrector,
+ corrector_kwargs=corrector_kwargs,
+ unconditional_guidance_scale=unconditional_guidance_scale,
+ unconditional_conditioning=unconditional_conditioning)
+ img, pred_x0 = outs
+ if callback: callback(i)
+ if img_callback: img_callback(pred_x0, i)
+
+ if index % log_every_t == 0 or index == total_steps - 1:
+ intermediates['x_inter'].append(img)
+ intermediates['pred_x0'].append(pred_x0)
+
+ return img, intermediates
+
+ @torch.no_grad()
+ def p_sample_ddim(self, x, c, t, index, repeat_noise=False, use_original_steps=False, quantize_denoised=False,
+ temperature=1., noise_dropout=0., score_corrector=None, corrector_kwargs=None,
+ unconditional_guidance_scale=1., unconditional_conditioning=None):
+ b, *_, device = *x.shape, x.device
+
+ if unconditional_conditioning is None or unconditional_guidance_scale == 1.:
+ e_t = self.model.apply_model(x, t, c)
+ else:
+ x_in = torch.cat([x] * 2)
+ t_in = torch.cat([t] * 2)
+ c_in = torch.cat([unconditional_conditioning, c])
+ e_t_uncond, e_t = self.model.apply_model(x_in, t_in, c_in).chunk(2)
+ e_t = e_t_uncond + unconditional_guidance_scale * (e_t - e_t_uncond)
+
+ if score_corrector is not None:
+ assert self.model.parameterization == "eps"
+ e_t = score_corrector.modify_score(self.model, e_t, x, t, c, **corrector_kwargs)
+
+ alphas = self.model.alphas_cumprod if use_original_steps else self.ddim_alphas
+ alphas_prev = self.model.alphas_cumprod_prev if use_original_steps else self.ddim_alphas_prev
+ sqrt_one_minus_alphas = self.model.sqrt_one_minus_alphas_cumprod if use_original_steps else self.ddim_sqrt_one_minus_alphas
+ sigmas = self.model.ddim_sigmas_for_original_num_steps if use_original_steps else self.ddim_sigmas
+ # select parameters corresponding to the currently considered timestep
+ a_t = torch.full((b, 1, 1, 1), alphas[index], device=device)
+ a_prev = torch.full((b, 1, 1, 1), alphas_prev[index], device=device)
+ sigma_t = torch.full((b, 1, 1, 1), sigmas[index], device=device)
+ sqrt_one_minus_at = torch.full((b, 1, 1, 1), sqrt_one_minus_alphas[index], device=device)
+
+ # current prediction for x_0
+ pred_x0 = (x - sqrt_one_minus_at * e_t) / a_t.sqrt()
+ if quantize_denoised:
+ pred_x0, _, *_ = self.model.first_stage_model.quantize(pred_x0)
+ # direction pointing to x_t
+ dir_xt = (1. - a_prev - sigma_t ** 2).sqrt() * e_t
+ noise = sigma_t * noise_like(x.shape, device, repeat_noise) * temperature
+ if noise_dropout > 0.:
+ noise = torch.nn.functional.dropout(noise, p=noise_dropout)
+ x_prev = a_prev.sqrt() * pred_x0 + dir_xt + noise
+ return x_prev, pred_x0
diff --git a/lidm/models/diffusion/ddpm.py b/lidm/models/diffusion/ddpm.py
new file mode 100644
index 0000000000000000000000000000000000000000..706e1c14188e40c3edc11d123ea74768e237cf1e
--- /dev/null
+++ b/lidm/models/diffusion/ddpm.py
@@ -0,0 +1,1455 @@
+"""
+wild mixture of
+https://github.com/lucidrains/denoising-diffusion-pytorch/blob/7706bdfc6f527f58d33f84b7b522e61e6e3164b3/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py
+https://github.com/openai/improved-diffusion/blob/e94489283bb876ac1477d5dd7709bbbd2d9902ce/improved_diffusion/gaussian_diffusion.py
+https://github.com/CompVis/taming-transformers
+-- merci
+"""
+
+import torch
+import torch.nn as nn
+import numpy as np
+import pytorch_lightning as pl
+from torch.optim.lr_scheduler import LambdaLR
+from einops import rearrange, repeat
+from contextlib import contextmanager
+from functools import partial
+from tqdm import tqdm
+from torchvision.utils import make_grid
+from pytorch_lightning.utilities.distributed import rank_zero_only
+
+from ...utils.misc_utils import log_txt_as_img, exists, default, ismap, isimage, mean_flat, count_params, instantiate_from_config, print_fn
+from ...modules.ema import LitEma
+from ...modules.distributions.distributions import normal_kl, DiagonalGaussianDistribution
+from ...models.autoencoder import VQModelInterface, IdentityFirstStage, AutoencoderKL
+from ...modules.basic import make_beta_schedule, extract_into_tensor, noise_like
+from ...models.diffusion.ddim import DDIMSampler
+
+__conditioning_keys__ = {'concat': 'c_concat',
+ 'crossattn': 'c_crossattn',
+ 'adm': 'y'}
+
+
+def disabled_train(self, mode=True):
+ """Overwrite model.train with this function to make sure train/eval mode
+ does not change anymore."""
+ return self
+
+
+def uniform_on_device(r1, r2, shape, device):
+ return (r1 - r2) * torch.rand(*shape, device=device) + r2
+
+
+class DDPM(pl.LightningModule):
+ # classic DDPM with Gaussian diffusion, in image space
+ def __init__(self,
+ unet_config,
+ timesteps=1000,
+ beta_schedule="linear",
+ loss_type="l2",
+ ckpt_path=None,
+ ignore_keys=[],
+ load_only_unet=False,
+ monitor="val/loss",
+ use_ema=True,
+ first_stage_key="image",
+ image_size=[256, 256],
+ channels=3,
+ log_every_t=100,
+ clip_denoised=True,
+ linear_start=1e-4,
+ linear_end=2e-2,
+ cosine_s=8e-3,
+ given_betas=None,
+ original_elbo_weight=0.,
+ v_posterior=0., # weight for choosing posterior variance as sigma = (1-v) * beta_tilde + v * beta
+ l_simple_weight=1.,
+ conditioning_key=None,
+ parameterization="eps", # all assuming fixed variance schedules
+ scheduler_config=None,
+ use_positional_encodings=False,
+ learn_logvar=False,
+ logvar_init=0.,
+ verbose=False
+ ):
+ super().__init__()
+ assert parameterization in ["eps", "x0"], 'currently only supporting "eps" and "x0"'
+ self.print_fn = partial(print_fn, verbose=verbose)
+ self.verbose = verbose
+ self.parameterization = parameterization
+ self.print_fn(f"{self.__class__.__name__}: Running in {self.parameterization}-prediction mode")
+ self.cond_stage_model = None
+ self.clip_denoised = clip_denoised
+ self.log_every_t = log_every_t
+ self.first_stage_key = first_stage_key
+ self.image_size = image_size # try conv?
+ self.channels = channels
+ self.use_positional_encodings = use_positional_encodings
+ self.model = DiffusionWrapper(unet_config, conditioning_key)
+ count_params(self.model, verbose=True)
+ self.use_ema = use_ema
+ if self.use_ema:
+ self.model_ema = LitEma(self.model)
+ self.print_fn(f"Keeping EMAs of {len(list(self.model_ema.buffers()))}.")
+
+ self.use_scheduler = scheduler_config is not None
+ if self.use_scheduler:
+ self.scheduler_config = scheduler_config
+
+ self.v_posterior = v_posterior
+ self.original_elbo_weight = original_elbo_weight
+ self.l_simple_weight = l_simple_weight
+
+ if monitor is not None:
+ self.monitor = monitor
+ if ckpt_path is not None:
+ self.init_from_ckpt(ckpt_path, ignore_keys=ignore_keys, only_model=load_only_unet)
+
+ self.register_schedule(given_betas=given_betas, beta_schedule=beta_schedule, timesteps=timesteps,
+ linear_start=linear_start, linear_end=linear_end, cosine_s=cosine_s)
+
+ self.loss_type = loss_type
+
+ self.learn_logvar = learn_logvar
+ self.logvar = torch.full(fill_value=logvar_init, size=(self.num_timesteps,))
+ self.logvar = nn.Parameter(self.logvar, requires_grad=self.learn_logvar)
+
+ def register_schedule(self, given_betas=None, beta_schedule="linear", timesteps=1000,
+ linear_start=1e-4, linear_end=2e-2, cosine_s=8e-3):
+ if exists(given_betas):
+ betas = given_betas
+ else:
+ betas = make_beta_schedule(beta_schedule, timesteps, linear_start=linear_start, linear_end=linear_end,
+ cosine_s=cosine_s)
+ alphas = 1. - betas
+ alphas_cumprod = np.cumprod(alphas, axis=0)
+ alphas_cumprod_prev = np.append(1., alphas_cumprod[:-1])
+
+ timesteps, = betas.shape
+ self.num_timesteps = int(timesteps)
+ self.linear_start = linear_start
+ self.linear_end = linear_end
+ assert alphas_cumprod.shape[0] == self.num_timesteps, 'alphas have to be defined for each timestep'
+
+ to_torch = partial(torch.tensor, dtype=torch.float32)
+
+ self.register_buffer('betas', to_torch(betas))
+ self.register_buffer('alphas_cumprod', to_torch(alphas_cumprod))
+ self.register_buffer('alphas_cumprod_prev', to_torch(alphas_cumprod_prev))
+
+ # calculations for diffusion q(x_t | x_{t-1}) and others
+ self.register_buffer('sqrt_alphas_cumprod', to_torch(np.sqrt(alphas_cumprod)))
+ self.register_buffer('sqrt_one_minus_alphas_cumprod', to_torch(np.sqrt(1. - alphas_cumprod)))
+ self.register_buffer('log_one_minus_alphas_cumprod', to_torch(np.log(1. - alphas_cumprod)))
+ self.register_buffer('sqrt_recip_alphas_cumprod', to_torch(np.sqrt(1. / alphas_cumprod)))
+ self.register_buffer('sqrt_recipm1_alphas_cumprod', to_torch(np.sqrt(1. / alphas_cumprod - 1)))
+
+ # calculations for posterior q(x_{t-1} | x_t, x_0)
+ posterior_variance = (1 - self.v_posterior) * betas * (1. - alphas_cumprod_prev) / (
+ 1. - alphas_cumprod) + self.v_posterior * betas
+ # above: equal to 1. / (1. / (1. - alpha_cumprod_tm1) + alpha_t / beta_t)
+ self.register_buffer('posterior_variance', to_torch(posterior_variance))
+ # below: log calculation clipped because the posterior variance is 0 at the beginning of the diffusion chain
+ self.register_buffer('posterior_log_variance_clipped', to_torch(np.log(np.maximum(posterior_variance, 1e-20))))
+ self.register_buffer('posterior_mean_coef1', to_torch(
+ betas * np.sqrt(alphas_cumprod_prev) / (1. - alphas_cumprod)))
+ self.register_buffer('posterior_mean_coef2', to_torch(
+ (1. - alphas_cumprod_prev) * np.sqrt(alphas) / (1. - alphas_cumprod)))
+
+ if self.parameterization == "eps":
+ lvlb_weights = self.betas ** 2 / (
+ 2 * self.posterior_variance * to_torch(alphas) * (1 - self.alphas_cumprod))
+ elif self.parameterization == "x0":
+ lvlb_weights = 0.5 * np.sqrt(torch.Tensor(alphas_cumprod)) / (2. * 1 - torch.Tensor(alphas_cumprod))
+ else:
+ raise NotImplementedError("mu not supported")
+
+ lvlb_weights[0] = lvlb_weights[1]
+ self.register_buffer('lvlb_weights', lvlb_weights, persistent=False)
+ assert not torch.isnan(self.lvlb_weights).all()
+
+ @contextmanager
+ def ema_scope(self, context=None):
+ if self.use_ema:
+ self.model_ema.store(self.model.parameters())
+ self.model_ema.copy_to(self.model)
+ if context is not None:
+ self.print_fn(f"{context}: Switched to EMA weights")
+ try:
+ yield None
+ finally:
+ if self.use_ema:
+ self.model_ema.restore(self.model.parameters())
+ if context is not None:
+ self.print_fn(f"{context}: Restored training weights")
+
+ def init_from_ckpt(self, path, ignore_keys=list(), only_model=False):
+ sd = torch.load(path, map_location="cpu")
+ if "state_dict" in list(sd.keys()):
+ sd = sd["state_dict"]
+ keys = list(sd.keys())
+ for k in keys:
+ for ik in ignore_keys:
+ if k.startswith(ik):
+ self.print_fn("Deleting key {} from state_dict.".format(k))
+ del sd[k]
+ missing, unexpected = self.load_state_dict(sd, strict=False) if not only_model else self.model.load_state_dict(
+ sd, strict=False)
+ self.print_fn(f"Restored from {path} with {len(missing)} missing and {len(unexpected)} unexpected keys")
+ if len(missing) > 0:
+ self.print_fn(f"Missing Keys: {missing}")
+ if len(unexpected) > 0:
+ self.print_fn(f"Unexpected Keys: {unexpected}")
+
+ def q_mean_variance(self, x_start, t):
+ """
+ Get the distribution q(x_t | x_0).
+ :param x_start: the [N x C x ...] tensor of noiseless inputs.
+ :param t: the number of diffusion steps (minus 1). Here, 0 means one step.
+ :return: A tuple (mean, variance, log_variance), all of x_start's shape.
+ """
+ mean = (extract_into_tensor(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start)
+ variance = extract_into_tensor(1.0 - self.alphas_cumprod, t, x_start.shape)
+ log_variance = extract_into_tensor(self.log_one_minus_alphas_cumprod, t, x_start.shape)
+ return mean, variance, log_variance
+
+ def predict_start_from_noise(self, x_t, t, noise):
+ return (
+ extract_into_tensor(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t -
+ extract_into_tensor(self.sqrt_recipm1_alphas_cumprod, t, x_t.shape) * noise
+ )
+
+ def q_posterior(self, x_start, x_t, t):
+ posterior_mean = (
+ extract_into_tensor(self.posterior_mean_coef1, t, x_t.shape) * x_start +
+ extract_into_tensor(self.posterior_mean_coef2, t, x_t.shape) * x_t
+ )
+ posterior_variance = extract_into_tensor(self.posterior_variance, t, x_t.shape)
+ posterior_log_variance_clipped = extract_into_tensor(self.posterior_log_variance_clipped, t, x_t.shape)
+ return posterior_mean, posterior_variance, posterior_log_variance_clipped
+
+ def p_mean_variance(self, x, t, clip_denoised: bool):
+ model_out = self.model(x, t)
+ if self.parameterization == "eps":
+ x_recon = self.predict_start_from_noise(x, t=t, noise=model_out)
+ elif self.parameterization == "x0":
+ x_recon = model_out
+ if clip_denoised:
+ x_recon.clamp_(-1., 1.)
+
+ model_mean, posterior_variance, posterior_log_variance = self.q_posterior(x_start=x_recon, x_t=x, t=t)
+ return model_mean, posterior_variance, posterior_log_variance
+
+ @torch.no_grad()
+ def p_sample(self, x, t, clip_denoised=True, repeat_noise=False):
+ b, *_, device = *x.shape, x.device
+ model_mean, _, model_log_variance = self.p_mean_variance(x=x, t=t, clip_denoised=clip_denoised)
+ noise = noise_like(x.shape, device, repeat_noise)
+ # no noise when t == 0
+ nonzero_mask = (1 - (t == 0).float()).reshape(b, *((1,) * (len(x.shape) - 1)))
+ return model_mean + nonzero_mask * (0.5 * model_log_variance).exp() * noise
+
+ @torch.no_grad()
+ def p_sample_loop(self, shape, return_intermediates=False):
+ device = self.betas.device
+ b = shape[0]
+ img = torch.randn(shape, device=device)
+ intermediates = [img]
+ for i in tqdm(reversed(range(0, self.num_timesteps)), desc='Sampling t', total=self.num_timesteps):
+ img = self.p_sample(img, torch.full((b,), i, device=device, dtype=torch.long),
+ clip_denoised=self.clip_denoised)
+ if i % self.log_every_t == 0 or i == self.num_timesteps - 1:
+ intermediates.append(img)
+ if return_intermediates:
+ return img, intermediates
+ return img
+
+ @torch.no_grad()
+ def sample(self, batch_size=16, return_intermediates=False):
+ image_size = self.image_size
+ channels = self.channels
+ return self.p_sample_loop((batch_size, channels, *image_size),
+ return_intermediates=return_intermediates)
+
+ def q_sample(self, x_start, t, noise=None):
+ noise = default(noise, lambda: torch.randn_like(x_start))
+ return (extract_into_tensor(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start +
+ extract_into_tensor(self.sqrt_one_minus_alphas_cumprod, t, x_start.shape) * noise)
+
+ def get_loss(self, pred, target, mean=True):
+ if self.loss_type == 'l1':
+ loss = (target - pred).abs()
+ if mean:
+ loss = loss.mean()
+ elif self.loss_type == 'l2':
+ if mean:
+ loss = torch.nn.functional.mse_loss(target, pred)
+ else:
+ loss = torch.nn.functional.mse_loss(target, pred, reduction='none')
+ else:
+ raise NotImplementedError("unknown loss type '{loss_type}'")
+
+ return loss
+
+ def p_losses(self, x_start, t, noise=None):
+ noise = default(noise, lambda: torch.randn_like(x_start))
+ x_noisy = self.q_sample(x_start=x_start, t=t, noise=noise)
+ model_out = self.model(x_noisy, t)
+
+ loss_dict = {}
+ if self.parameterization == "eps":
+ target = noise
+ elif self.parameterization == "x0":
+ target = x_start
+ else:
+ raise NotImplementedError(f"Paramterization {self.parameterization} not yet supported")
+
+ loss = self.get_loss(model_out, target, mean=False).mean(dim=[1, 2, 3])
+
+ log_prefix = 'train' if self.training else 'val'
+
+ loss_dict.update({f'{log_prefix}/loss_simple': loss.mean()})
+ loss_simple = loss.mean() * self.l_simple_weight
+
+ loss_vlb = (self.lvlb_weights[t] * loss).mean()
+ loss_dict.update({f'{log_prefix}/loss_vlb': loss_vlb})
+
+ loss = loss_simple + self.original_elbo_weight * loss_vlb
+
+ loss_dict.update({f'{log_prefix}/loss': loss})
+
+ return loss, loss_dict
+
+ def forward(self, x, *args, **kwargs):
+ # b, c, h, w, device, img_size, = *x.shape, x.device, self.image_size
+ # assert h == img_size and w == img_size, f'height and width of image must be {img_size}'
+ t = torch.randint(0, self.num_timesteps, (x.shape[0],), device=self.device).long()
+ return self.p_losses(x, t, *args, **kwargs)
+
+ def get_input(self, batch, k):
+ x = batch[k]
+ # if len(x.shape) == 3:
+ # x = x[..., None]
+ return x
+
+ def shared_step(self, batch):
+ x = self.get_input(batch, self.first_stage_key)
+ loss, loss_dict = self(x)
+ return loss, loss_dict
+
+ def training_step(self, batch, batch_idx):
+ loss, loss_dict = self.shared_step(batch)
+
+ self.log_dict(loss_dict, prog_bar=True,
+ logger=True, on_step=True, on_epoch=True)
+
+ self.log("global_step", self.global_step,
+ prog_bar=True, logger=True, on_step=True, on_epoch=False)
+
+ if self.use_scheduler:
+ lr = self.optimizers().param_groups[0]['lr']
+ self.log('lr_abs', lr, prog_bar=True, logger=True, on_step=True, on_epoch=False)
+
+ return loss
+
+ @torch.no_grad()
+ def validation_step(self, batch, batch_idx):
+ _, loss_dict_no_ema = self.shared_step(batch)
+ with self.ema_scope():
+ _, loss_dict_ema = self.shared_step(batch)
+ loss_dict_ema = {key + '_ema': loss_dict_ema[key] for key in loss_dict_ema}
+ self.log_dict(loss_dict_no_ema, prog_bar=False, logger=True, on_step=False, on_epoch=True)
+ self.log_dict(loss_dict_ema, prog_bar=False, logger=True, on_step=False, on_epoch=True)
+
+ def on_train_batch_end(self, *args, **kwargs):
+ if self.use_ema:
+ self.model_ema(self.model)
+
+ def _get_rows_from_list(self, samples):
+ n_imgs_per_row = len(samples)
+ denoise_grid = rearrange(samples, 'n b c h w -> b n c h w')
+ denoise_grid = rearrange(denoise_grid, 'b n c h w -> (b n) c h w')
+ denoise_grid = make_grid(denoise_grid, nrow=n_imgs_per_row)
+ return denoise_grid
+
+ @torch.no_grad()
+ def log_images(self, batch, N=8, n_row=2, sample=True, return_keys=None, **kwargs):
+ log = dict()
+ x = self.get_input(batch, self.first_stage_key)
+ N = min(x.shape[0], N)
+ n_row = min(x.shape[0], n_row)
+ x = x.to(self.device)[:N]
+ log["inputs"] = x
+
+ # get diffusion row
+ diffusion_row = list()
+ x_start = x[:n_row]
+
+ for t in range(self.num_timesteps):
+ if t % self.log_every_t == 0 or t == self.num_timesteps - 1:
+ t = repeat(torch.tensor([t]), '1 -> b', b=n_row)
+ t = t.to(self.device).long()
+ noise = torch.randn_like(x_start)
+ x_noisy = self.q_sample(x_start=x_start, t=t, noise=noise)
+ diffusion_row.append(x_noisy)
+
+ log["diffusion_row"] = self._get_rows_from_list(diffusion_row)
+
+ if sample:
+ # get denoise row
+ with self.ema_scope("Plotting"):
+ samples, denoise_row = self.sample(batch_size=N, return_intermediates=True)
+
+ log["samples"] = samples
+ log["denoise_row"] = self._get_rows_from_list(denoise_row)
+
+ if return_keys:
+ if np.intersect1d(list(log.keys()), return_keys).shape[0] == 0:
+ return log
+ else:
+ return {key: log[key] for key in return_keys}
+ return log
+
+ def configure_optimizers(self):
+ lr = self.learning_rate
+ params = list(self.model.parameters())
+ if self.learn_logvar:
+ params = params + [self.logvar]
+ opt = torch.optim.AdamW(params, lr=lr)
+ return opt
+
+
+class LatentDiffusion(DDPM):
+ """main class"""
+
+ def __init__(self,
+ first_stage_config,
+ cond_stage_config,
+ num_timesteps_cond=None,
+ cond_stage_key="image",
+ cond_stage_trainable=False,
+ concat_mode=True,
+ cond_stage_forward=None,
+ conditioning_key=None,
+ scale_factor=1.0,
+ scale_by_std=False,
+ use_mask=False,
+ *args, **kwargs):
+ self.num_timesteps_cond = default(num_timesteps_cond, 1)
+ self.scale_by_std = scale_by_std
+ assert self.num_timesteps_cond <= kwargs['timesteps']
+ # for backwards compatibility after implementation of DiffusionWrapper
+ if conditioning_key is None:
+ conditioning_key = 'concat' if concat_mode else 'crossattn'
+ if cond_stage_config == '__is_unconditional__':
+ conditioning_key = None
+ ckpt_path = kwargs.pop("ckpt_path", None)
+ ignore_keys = kwargs.pop("ignore_keys", [])
+ super().__init__(conditioning_key=conditioning_key, *args, **kwargs)
+ self.concat_mode = concat_mode
+ self.cond_stage_trainable = cond_stage_trainable
+ self.cond_stage_key = cond_stage_key
+ try:
+ self.num_downs = len(first_stage_config.params.ddconfig.ch_mult) - 1
+ except:
+ self.num_downs = 0
+ if not scale_by_std:
+ self.scale_factor = scale_factor
+ else:
+ self.register_buffer('scale_factor', torch.tensor(scale_factor))
+ self.instantiate_first_stage(first_stage_config)
+ self.instantiate_cond_stage(cond_stage_config)
+ self.cond_stage_forward = cond_stage_forward
+ self.clip_denoised = False
+ self.bbox_tokenizer = None
+ self.use_mask = use_mask
+
+ self.restarted_from_ckpt = False
+ if ckpt_path is not None:
+ self.init_from_ckpt(ckpt_path, ignore_keys)
+ self.restarted_from_ckpt = True
+
+ def make_cond_schedule(self, ):
+ self.cond_ids = torch.full(size=(self.num_timesteps,), fill_value=self.num_timesteps - 1, dtype=torch.long)
+ ids = torch.round(torch.linspace(0, self.num_timesteps - 1, self.num_timesteps_cond)).long()
+ self.cond_ids[:self.num_timesteps_cond] = ids
+
+ @rank_zero_only
+ @torch.no_grad()
+ def on_train_batch_start(self, batch, batch_idx, dataloader_idx):
+ # only for very first batch
+ if self.scale_by_std and self.current_epoch == 0 and self.global_step == 0 and batch_idx == 0 and not self.restarted_from_ckpt:
+ assert self.scale_factor == 1., 'rather not use custom rescaling and std-rescaling simultaneously'
+ # set rescale weight to 1./std of encodings
+ self.print_fn("### USING STD-RESCALING ###")
+ x = super().get_input(batch, self.first_stage_key)
+ x = x.to(self.device)
+ encoder_posterior = self.encode_first_stage(x)
+ z = self.get_first_stage_encoding(encoder_posterior).detach()
+ del self.scale_factor
+ self.register_buffer('scale_factor', 1. / z.flatten().std())
+ self.print_fn(f"setting self.scale_factor to {self.scale_factor}")
+ self.print_fn("### USING STD-RESCALING ###")
+
+ def register_schedule(self,
+ given_betas=None, beta_schedule="linear", timesteps=1000,
+ linear_start=1e-4, linear_end=2e-2, cosine_s=8e-3):
+ super().register_schedule(given_betas, beta_schedule, timesteps, linear_start, linear_end, cosine_s)
+
+ self.shorten_cond_schedule = self.num_timesteps_cond > 1
+ if self.shorten_cond_schedule:
+ self.make_cond_schedule()
+
+ def instantiate_first_stage(self, config):
+ model = instantiate_from_config(config)
+ self.first_stage_model = model.eval()
+ self.first_stage_model.train = disabled_train
+ for param in self.first_stage_model.parameters():
+ param.requires_grad = False
+
+ def instantiate_cond_stage(self, config):
+ if not self.cond_stage_trainable:
+ if config == "__is_first_stage__":
+ self.print_fn("Using first stage also as cond stage.")
+ self.cond_stage_model = self.first_stage_model
+ elif config == "__is_unconditional__":
+ self.print_fn(f"Training {self.__class__.__name__} as an unconditional model.")
+ self.cond_stage_model = None
+ # self.be_unconditional = True
+ else:
+ model = instantiate_from_config(config)
+ self.cond_stage_model = model.eval()
+ self.cond_stage_model.train = disabled_train
+ for param in self.cond_stage_model.parameters():
+ param.requires_grad = False
+ else:
+ assert config != '__is_first_stage__'
+ assert config != '__is_unconditional__'
+ model = instantiate_from_config(config)
+ self.cond_stage_model = model
+
+ def _get_denoise_row_from_list(self, samples, desc='', force_no_decoder_quantization=False):
+ denoise_row = []
+ for zd in tqdm(samples, desc=desc):
+ denoise_row.append(self.decode_first_stage(zd.to(self.device),
+ force_not_quantize=force_no_decoder_quantization))
+ n_imgs_per_row = len(denoise_row)
+ denoise_row = torch.stack(denoise_row) # n_log_step, n_row, C, H, W
+ denoise_grid = rearrange(denoise_row, 'n b c h w -> b n c h w')
+ denoise_grid = rearrange(denoise_grid, 'b n c h w -> (b n) c h w')
+ denoise_grid = make_grid(denoise_grid, nrow=n_imgs_per_row)
+ return denoise_grid
+
+ def get_first_stage_encoding(self, encoder_posterior):
+ if isinstance(encoder_posterior, DiagonalGaussianDistribution):
+ z = encoder_posterior.sample()
+ elif isinstance(encoder_posterior, torch.Tensor):
+ z = encoder_posterior
+ else:
+ raise NotImplementedError(f"encoder_posterior of type '{type(encoder_posterior)}' not yet implemented")
+ return self.scale_factor * z
+
+ def get_learned_conditioning(self, c):
+ if self.cond_stage_forward is None:
+ if hasattr(self.cond_stage_model, 'encode') and callable(self.cond_stage_model.encode):
+ c = self.cond_stage_model.encode(c)
+ if isinstance(c, DiagonalGaussianDistribution):
+ c = c.mode()
+ else:
+ c = self.cond_stage_model(c)
+ else:
+ assert hasattr(self.cond_stage_model, self.cond_stage_forward)
+ c = getattr(self.cond_stage_model, self.cond_stage_forward)(c)
+ return c
+
+ def meshgrid(self, h, w):
+ y = torch.arange(0, h).view(h, 1, 1).repeat(1, w, 1)
+ x = torch.arange(0, w).view(1, w, 1).repeat(h, 1, 1)
+
+ arr = torch.cat([y, x], dim=-1)
+ return arr
+
+ def delta_border(self, h, w):
+ """
+ :param h: height
+ :param w: width
+ :return: normalized distance to image border,
+ wtith min distance = 0 at border and max dist = 0.5 at image center
+ """
+ lower_right_corner = torch.tensor([h - 1, w - 1]).view(1, 1, 2)
+ arr = self.meshgrid(h, w) / lower_right_corner
+ dist_left_up = torch.min(arr, dim=-1, keepdims=True)[0]
+ dist_right_down = torch.min(1 - arr, dim=-1, keepdims=True)[0]
+ edge_dist = torch.min(torch.cat([dist_left_up, dist_right_down], dim=-1), dim=-1)[0]
+ return edge_dist
+
+ def get_weighting(self, h, w, Ly, Lx, device):
+ weighting = self.delta_border(h, w)
+ weighting = torch.clip(weighting, self.split_input_params["clip_min_weight"],
+ self.split_input_params["clip_max_weight"], )
+ weighting = weighting.view(1, h * w, 1).repeat(1, 1, Ly * Lx).to(device)
+
+ if self.split_input_params["tie_braker"]:
+ L_weighting = self.delta_border(Ly, Lx)
+ L_weighting = torch.clip(L_weighting,
+ self.split_input_params["clip_min_tie_weight"],
+ self.split_input_params["clip_max_tie_weight"])
+
+ L_weighting = L_weighting.view(1, 1, Ly * Lx).to(device)
+ weighting = weighting * L_weighting
+ return weighting
+
+ def get_fold_unfold(self, x, kernel_size, stride, uf=1, df=1): # todo load once not every time, shorten code
+ """
+ :param x: img of size (bs, c, h, w)
+ :return: n img crops of size (n, bs, c, kernel_size[0], kernel_size[1])
+ """
+ bs, nc, h, w = x.shape
+
+ # number of crops in image
+ Ly = (h - kernel_size[0]) // stride[0] + 1
+ Lx = (w - kernel_size[1]) // stride[1] + 1
+
+ if uf == 1 and df == 1:
+ fold_params = dict(kernel_size=kernel_size, dilation=1, padding=0, stride=stride)
+ unfold = torch.nn.Unfold(**fold_params)
+
+ fold = torch.nn.Fold(output_size=x.shape[2:], **fold_params)
+
+ weighting = self.get_weighting(kernel_size[0], kernel_size[1], Ly, Lx, x.device).to(x.dtype)
+ normalization = fold(weighting).view(1, 1, h, w) # normalizes the overlap
+ weighting = weighting.view((1, 1, kernel_size[0], kernel_size[1], Ly * Lx))
+
+ elif uf > 1 and df == 1:
+ fold_params = dict(kernel_size=kernel_size, dilation=1, padding=0, stride=stride)
+ unfold = torch.nn.Unfold(**fold_params)
+
+ fold_params2 = dict(kernel_size=(kernel_size[0] * uf, kernel_size[0] * uf),
+ dilation=1, padding=0,
+ stride=(stride[0] * uf, stride[1] * uf))
+ fold = torch.nn.Fold(output_size=(x.shape[2] * uf, x.shape[3] * uf), **fold_params2)
+
+ weighting = self.get_weighting(kernel_size[0] * uf, kernel_size[1] * uf, Ly, Lx, x.device).to(x.dtype)
+ normalization = fold(weighting).view(1, 1, h * uf, w * uf) # normalizes the overlap
+ weighting = weighting.view((1, 1, kernel_size[0] * uf, kernel_size[1] * uf, Ly * Lx))
+
+ elif df > 1 and uf == 1:
+ fold_params = dict(kernel_size=kernel_size, dilation=1, padding=0, stride=stride)
+ unfold = torch.nn.Unfold(**fold_params)
+
+ fold_params2 = dict(kernel_size=(kernel_size[0] // df, kernel_size[0] // df),
+ dilation=1, padding=0,
+ stride=(stride[0] // df, stride[1] // df))
+ fold = torch.nn.Fold(output_size=(x.shape[2] // df, x.shape[3] // df), **fold_params2)
+
+ weighting = self.get_weighting(kernel_size[0] // df, kernel_size[1] // df, Ly, Lx, x.device).to(x.dtype)
+ normalization = fold(weighting).view(1, 1, h // df, w // df) # normalizes the overlap
+ weighting = weighting.view((1, 1, kernel_size[0] // df, kernel_size[1] // df, Ly * Lx))
+
+ else:
+ raise NotImplementedError
+
+ return fold, unfold, normalization, weighting
+
+ @torch.no_grad()
+ def get_input(self, batch, k, return_first_stage_outputs=False, force_c_encode=False,
+ cond_key=None, return_original_cond=False, bs=None):
+ # ground truth
+ x = super().get_input(batch, k)
+ if bs is not None:
+ x = x[:bs]
+ x = x.to(self.device)
+
+ # encoding
+ encoder_posterior = self.encode_first_stage(x)
+ z = self.get_first_stage_encoding(encoder_posterior).detach()
+ if self.model.conditioning_key is not None:
+ if cond_key is None:
+ cond_key = self.cond_stage_key
+ if cond_key != self.first_stage_key:
+ if cond_key in ['caption', 'bbox', 'center', 'camera']:
+ xc = batch[cond_key]
+ elif cond_key in ['class_label']:
+ xc = batch
+ else:
+ xc = super().get_input(batch, cond_key).to(self.device)
+ else:
+ xc = x
+ # if bs is not None:
+ # xc = xc[:bs]
+ if not self.cond_stage_trainable or force_c_encode:
+ if isinstance(xc, (dict, list)):
+ c = self.get_learned_conditioning(xc)
+ else:
+ c = self.get_learned_conditioning(xc.to(self.device))
+ else:
+ c = xc
+ if bs is not None:
+ c = c[:bs]
+
+ if self.use_positional_encodings:
+ pos_x, pos_y = self.compute_latent_shifts(batch)
+ ckey = __conditioning_keys__[self.model.conditioning_key]
+ c = {ckey: c, 'pos_x': pos_x, 'pos_y': pos_y}
+
+ else:
+ c = None
+ xc = None
+ if self.use_positional_encodings:
+ pos_x, pos_y = self.compute_latent_shifts(batch)
+ c = {'pos_x': pos_x, 'pos_y': pos_y}
+ out = [z, c]
+ if return_first_stage_outputs:
+ xrec = self.decode_first_stage(z)
+ out.extend([x, xrec])
+ if return_original_cond:
+ out.append(xc)
+ return out
+
+ @torch.no_grad()
+ def decode_first_stage(self, z, predict_cids=False, force_not_quantize=False):
+ if predict_cids:
+ if z.dim() == 4:
+ z = torch.argmax(z.exp(), dim=1).long()
+ z = self.first_stage_model.quantize.get_codebook_entry(z, shape=None)
+ z = rearrange(z, 'b h w c -> b c h w').contiguous()
+
+ z = 1. / self.scale_factor * z
+
+ if hasattr(self, "split_input_params"):
+ if self.split_input_params["patch_distributed_vq"]:
+ ks = self.split_input_params["ks"] # eg. (128, 128)
+ stride = self.split_input_params["stride"] # eg. (64, 64)
+ uf = self.split_input_params["vqf"]
+ bs, nc, h, w = z.shape
+ if ks[0] > h or ks[1] > w:
+ ks = (min(ks[0], h), min(ks[1], w))
+ self.print_fn("reducing Kernel")
+
+ if stride[0] > h or stride[1] > w:
+ stride = (min(stride[0], h), min(stride[1], w))
+ self.print_fn("reducing stride")
+
+ fold, unfold, normalization, weighting = self.get_fold_unfold(z, ks, stride, uf=uf)
+
+ z = unfold(z) # (bn, nc * prod(**ks), L)
+ # 1. Reshape to img shape
+ z = z.view((z.shape[0], -1, ks[0], ks[1], z.shape[-1])) # (bn, nc, ks[0], ks[1], L )
+
+ # 2. apply model loop over last dim
+ if isinstance(self.first_stage_model, VQModelInterface):
+ output_list = [self.first_stage_model.decode(z[:, :, :, :, i],
+ force_not_quantize=predict_cids or force_not_quantize)
+ for i in range(z.shape[-1])]
+ else:
+
+ output_list = [self.first_stage_model.decode(z[:, :, :, :, i])
+ for i in range(z.shape[-1])]
+
+ o = torch.stack(output_list, axis=-1) # # (bn, nc, ks[0], ks[1], L)
+ o = o * weighting
+ # Reverse 1. reshape to img shape
+ o = o.view((o.shape[0], -1, o.shape[-1])) # (bn, nc * ks[0] * ks[1], L)
+ # stitch crops together
+ decoded = fold(o)
+ decoded = decoded / normalization # norm is shape (1, 1, h, w)
+ return decoded
+ else:
+ if isinstance(self.first_stage_model, VQModelInterface):
+ return self.first_stage_model.decode(z, force_not_quantize=predict_cids or force_not_quantize)
+ else:
+ return self.first_stage_model.decode(z)
+
+ else:
+ if isinstance(self.first_stage_model, VQModelInterface):
+ return self.first_stage_model.decode(z, force_not_quantize=predict_cids or force_not_quantize)
+ else:
+ return self.first_stage_model.decode(z)
+
+ # same as above but without decorator
+ def differentiable_decode_first_stage(self, z, predict_cids=False, force_not_quantize=False):
+ if predict_cids:
+ if z.dim() == 4:
+ z = torch.argmax(z.exp(), dim=1).long()
+ z = self.first_stage_model.quantize.get_codebook_entry(z, shape=None)
+ z = rearrange(z, 'b h w c -> b c h w').contiguous()
+
+ z = 1. / self.scale_factor * z
+
+ if hasattr(self, "split_input_params"):
+ if self.split_input_params["patch_distributed_vq"]:
+ ks = self.split_input_params["ks"] # eg. (128, 128)
+ stride = self.split_input_params["stride"] # eg. (64, 64)
+ uf = self.split_input_params["vqf"]
+ bs, nc, h, w = z.shape
+ if ks[0] > h or ks[1] > w:
+ ks = (min(ks[0], h), min(ks[1], w))
+ self.print_fn("reducing Kernel")
+
+ if stride[0] > h or stride[1] > w:
+ stride = (min(stride[0], h), min(stride[1], w))
+ self.print_fn("reducing stride")
+
+ fold, unfold, normalization, weighting = self.get_fold_unfold(z, ks, stride, uf=uf)
+
+ z = unfold(z) # (bn, nc * prod(**ks), L)
+ # 1. Reshape to img shape
+ z = z.view((z.shape[0], -1, ks[0], ks[1], z.shape[-1])) # (bn, nc, ks[0], ks[1], L )
+
+ # 2. apply model loop over last dim
+ if isinstance(self.first_stage_model, VQModelInterface):
+ output_list = [self.first_stage_model.decode(z[:, :, :, :, i],
+ force_not_quantize=predict_cids or force_not_quantize)
+ for i in range(z.shape[-1])]
+ else:
+ output_list = [self.first_stage_model.decode(z[:, :, :, :, i])
+ for i in range(z.shape[-1])]
+
+ o = torch.stack(output_list, axis=-1) # # (bn, nc, ks[0], ks[1], L)
+ o = o * weighting
+ # Reverse 1. reshape to img shape
+ o = o.view((o.shape[0], -1, o.shape[-1])) # (bn, nc * ks[0] * ks[1], L)
+ # stitch crops together
+ decoded = fold(o)
+ decoded = decoded / normalization # norm is shape (1, 1, h, w)
+ return decoded
+ else:
+ if isinstance(self.first_stage_model, VQModelInterface):
+ return self.first_stage_model.decode(z, force_not_quantize=predict_cids or force_not_quantize)
+ else:
+ return self.first_stage_model.decode(z)
+
+ else:
+ if isinstance(self.first_stage_model, VQModelInterface):
+ return self.first_stage_model.decode(z, force_not_quantize=predict_cids or force_not_quantize)
+ else:
+ return self.first_stage_model.decode(z)
+
+ @torch.no_grad()
+ def encode_first_stage(self, x):
+ if hasattr(self, "split_input_params"):
+ if self.split_input_params["patch_distributed_vq"]:
+ ks = self.split_input_params["ks"] # eg. (128, 128)
+ stride = self.split_input_params["stride"] # eg. (64, 64)
+ df = self.split_input_params["vqf"]
+ self.split_input_params['original_image_size'] = x.shape[-2:]
+ bs, nc, h, w = x.shape
+ if ks[0] > h or ks[1] > w:
+ ks = (min(ks[0], h), min(ks[1], w))
+ self.print_fn("reducing kernel")
+
+ if stride[0] > h or stride[1] > w:
+ stride = (min(stride[0], h), min(stride[1], w))
+ self.print_fn("reducing stride")
+
+ fold, unfold, normalization, weighting = self.get_fold_unfold(x, ks, stride, df=df)
+ z = unfold(x) # (bn, nc * prod(**ks), L)
+ # Reshape to img shape
+ z = z.view((z.shape[0], -1, ks[0], ks[1], z.shape[-1])) # (bn, nc, ks[0], ks[1], L )
+
+ output_list = [self.first_stage_model.encode(z[:, :, :, :, i]) for i in range(z.shape[-1])]
+
+ o = torch.stack(output_list, axis=-1)
+ o = o * weighting
+
+ # Reverse reshape to img shape
+ o = o.view((o.shape[0], -1, o.shape[-1])) # (bn, nc * ks[0] * ks[1], L)
+ # stitch crops together
+ decoded = fold(o)
+ decoded = decoded / normalization
+ return decoded
+ else:
+ return self.first_stage_model.encode(x)
+ else:
+ return self.first_stage_model.encode(x)
+
+ def shared_step(self, batch, **kwargs):
+ x, c = self.get_input(batch, self.first_stage_key)
+ loss = self(x, c)
+ return loss
+
+ def forward(self, x, c, *args, **kwargs):
+ t = torch.randint(0, self.num_timesteps, (x.shape[0],), device=self.device).long()
+ if self.model.conditioning_key is not None:
+ assert c is not None
+ if self.cond_stage_trainable:
+ c = self.get_learned_conditioning(c)
+ if self.shorten_cond_schedule: # TODO: drop this option
+ tc = self.cond_ids[t].to(self.device)
+ c = self.q_sample(x_start=c, t=tc, noise=torch.randn_like(c.float()))
+ return self.p_losses(x, c, t, *args, **kwargs)
+
+ def _rescale_annotations(self, bboxes, crop_coordinates): # TODO: move to dataset
+ def rescale_bbox(bbox):
+ x0 = clamp((bbox[0] - crop_coordinates[0]) / crop_coordinates[2])
+ y0 = clamp((bbox[1] - crop_coordinates[1]) / crop_coordinates[3])
+ w = min(bbox[2] / crop_coordinates[2], 1 - x0)
+ h = min(bbox[3] / crop_coordinates[3], 1 - y0)
+ return x0, y0, w, h
+
+ return [rescale_bbox(b) for b in bboxes]
+
+ def apply_model(self, x_noisy, t, cond, return_ids=False):
+
+ if isinstance(cond, dict):
+ # hybrid case, cond is exptected to be a dict
+ pass
+ else:
+ if not isinstance(cond, list):
+ cond = [cond]
+ key = 'c_concat' if self.model.conditioning_key == 'concat' else 'c_crossattn'
+ cond = {key: cond}
+
+ if hasattr(self, "split_input_params"):
+ assert len(cond) == 1
+ assert not return_ids
+ ks = self.split_input_params["ks"] # eg. (128, 128)
+ stride = self.split_input_params["stride"] # eg. (64, 64)
+
+ h, w = x_noisy.shape[-2:]
+
+ fold, unfold, normalization, weighting = self.get_fold_unfold(x_noisy, ks, stride)
+
+ z = unfold(x_noisy) # (bn, nc * prod(**ks), L)
+ # Reshape to img shape
+ z = z.view((z.shape[0], -1, ks[0], ks[1], z.shape[-1])) # (bn, nc, ks[0], ks[1], L )
+ z_list = [z[:, :, :, :, i] for i in range(z.shape[-1])]
+
+ if self.cond_stage_key in ["image", "LR_image", "segmentation", 'bbox_img'] and self.model.conditioning_key:
+ c_key = next(iter(cond.keys())) # get key
+ c = next(iter(cond.values())) # get value
+ assert (len(c) == 1)
+ c = c[0] # get element
+
+ c = unfold(c)
+ c = c.view((c.shape[0], -1, ks[0], ks[1], c.shape[-1])) # (bn, nc, ks[0], ks[1], L )
+
+ cond_list = [{c_key: [c[:, :, :, :, i]]} for i in range(c.shape[-1])]
+
+ elif self.cond_stage_key in ['bbox', 'center']:
+ assert 'original_image_size' in self.split_input_params, 'BoudingBoxRescaling is missing original_image_size'
+
+ # assuming padding of unfold is always 0 and its dilation is always 1
+ n_patches_per_row = int((w - ks[0]) / stride[0] + 1)
+ full_img_h, full_img_w = self.split_input_params['original_image_size']
+ # as we are operating on latents, we need the factor from the original image size to the
+ # spatial latent size to properly rescale the crops for regenerating the bbox annotations
+ num_downs = self.first_stage_model.encoder.num_resolutions - 1
+ rescale_latent = 2 ** num_downs
+
+ # get top left postions of patches as conforming for the bbbox tokenizer, therefore we
+ # need to rescale the tl patch coordinates to be in between (0,1)
+ tl_patch_coordinates = [(rescale_latent * stride[0] * (patch_nr % n_patches_per_row) / full_img_w,
+ rescale_latent * stride[1] * (patch_nr // n_patches_per_row) / full_img_h)
+ for patch_nr in range(z.shape[-1])]
+
+ # patch_limits are tl_coord, width and height coordinates as (x_tl, y_tl, h, w)
+ patch_limits = [(x_tl, y_tl,
+ rescale_latent * ks[0] / full_img_w,
+ rescale_latent * ks[1] / full_img_h) for x_tl, y_tl in tl_patch_coordinates]
+ # patch_values = [(np.arange(x_tl,min(x_tl+ks, 1.)),np.arange(y_tl,min(y_tl+ks, 1.))) for x_tl, y_tl in tl_patch_coordinates]
+
+ # tokenize crop coordinates for the bounding boxes of the respective patches
+ patch_limits_tknzd = [torch.LongTensor(self.bbox_tokenizer.crop_encoder(bbox))[None].to(self.device)
+ for bbox in patch_limits] # list of length l with tensors of shape (1, 2)
+ self.print_fn(patch_limits_tknzd[0].shape)
+ # cut tknzd crop position from conditioning
+ assert isinstance(cond, dict), 'cond must be dict to be fed into model'
+ cut_cond = cond['c_crossattn'][0][..., :-2].to(self.device)
+ self.print_fn(cut_cond.shape)
+
+ adapted_cond = torch.stack([torch.cat([cut_cond, p], dim=1) for p in patch_limits_tknzd])
+ adapted_cond = rearrange(adapted_cond, 'l b n -> (l b) n')
+ self.print_fn(adapted_cond.shape)
+ adapted_cond = self.get_learned_conditioning(adapted_cond)
+ self.print_fn(adapted_cond.shape)
+ adapted_cond = rearrange(adapted_cond, '(l b) n d -> l b n d', l=z.shape[-1])
+ self.print_fn(adapted_cond.shape)
+
+ cond_list = [{'c_crossattn': [e]} for e in adapted_cond]
+
+ else:
+ cond_list = [cond for i in range(z.shape[-1])] # todo make this more efficient
+
+ # apply model by loop over crops
+ output_list = [self.model(z_list[i], t, **cond_list[i]) for i in range(z.shape[-1])]
+ assert not isinstance(output_list[0],
+ tuple) # todo cant deal with multiple model outputs check this never happens
+
+ o = torch.stack(output_list, axis=-1)
+ o = o * weighting
+ # Reverse reshape to img shape
+ o = o.view((o.shape[0], -1, o.shape[-1])) # (bn, nc * ks[0] * ks[1], L)
+ # stitch crops together
+ x_recon = fold(o) / normalization
+
+ else:
+ x_recon = self.model(x_noisy, t, **cond)
+
+ if isinstance(x_recon, tuple) and not return_ids:
+ return x_recon[0]
+ else:
+ return x_recon
+
+ def _predict_eps_from_xstart(self, x_t, t, pred_xstart):
+ return (extract_into_tensor(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t - pred_xstart) / \
+ extract_into_tensor(self.sqrt_recipm1_alphas_cumprod, t, x_t.shape)
+
+ def _prior_bpd(self, x_start):
+ """
+ Get the prior KL term for the variational lower-bound, measured in
+ bits-per-dim.
+ This term can't be optimized, as it only depends on the encoder.
+ :param x_start: the [N x C x ...] tensor of inputs.
+ :return: a batch of [N] KL values (in bits), one per batch element.
+ """
+ batch_size = x_start.shape[0]
+ t = torch.tensor([self.num_timesteps - 1] * batch_size, device=x_start.device)
+ qt_mean, _, qt_log_variance = self.q_mean_variance(x_start, t)
+ kl_prior = normal_kl(mean1=qt_mean, logvar1=qt_log_variance, mean2=0.0, logvar2=0.0)
+ return mean_flat(kl_prior) / np.log(2.0)
+
+ def p_losses(self, x_start, cond, t, noise=None):
+ noise = default(noise, lambda: torch.randn_like(x_start))
+ x_noisy = self.q_sample(x_start=x_start, t=t, noise=noise)
+ model_output = self.apply_model(x_noisy, t, cond)
+
+ loss_dict = {}
+ prefix = 'train' if self.training else 'val'
+
+ if self.parameterization == "x0":
+ target = x_start
+ elif self.parameterization == "eps":
+ target = noise
+ else:
+ raise NotImplementedError()
+
+ # simple loss
+ loss_simple = self.get_loss(model_output, target, mean=False).mean([1, 2, 3])
+ loss_dict.update({f'{prefix}/loss_simple': loss_simple.mean()})
+
+ logvar_t = self.logvar[t].to(self.device)
+ loss = loss_simple / torch.exp(logvar_t) + logvar_t
+ # loss = loss_simple / torch.exp(self.logvar) + self.logvar
+ if self.learn_logvar:
+ loss_dict.update({f'{prefix}/loss_gamma': loss.mean()})
+ loss_dict.update({'logvar': self.logvar.data.mean()})
+
+ loss = self.l_simple_weight * loss.mean()
+
+ # vlb loss
+ loss_vlb = self.get_loss(model_output, target, mean=False).mean(dim=(1, 2, 3))
+ loss_vlb = (self.lvlb_weights[t] * loss_vlb).mean()
+ loss_dict.update({f'{prefix}/loss_vlb': loss_vlb})
+
+ # total loss
+ loss += (self.original_elbo_weight * loss_vlb)
+ loss_dict.update({f'{prefix}/loss': loss})
+
+ return loss, loss_dict
+
+ def p_mean_variance(self, x, c, t, clip_denoised: bool, return_codebook_ids=False, quantize_denoised=False,
+ return_x0=False, score_corrector=None, corrector_kwargs=None):
+ t_in = t
+ model_out = self.apply_model(x, t_in, c, return_ids=return_codebook_ids)
+
+ if score_corrector is not None:
+ assert self.parameterization == "eps"
+ model_out = score_corrector.modify_score(self, model_out, x, t, c, **corrector_kwargs)
+
+ if return_codebook_ids:
+ model_out, logits = model_out
+
+ if self.parameterization == "eps":
+ x_recon = self.predict_start_from_noise(x, t=t, noise=model_out)
+ elif self.parameterization == "x0":
+ x_recon = model_out
+ else:
+ raise NotImplementedError()
+
+ if clip_denoised:
+ x_recon.clamp_(-1., 1.)
+ if quantize_denoised:
+ x_recon, _, [_, _, indices] = self.first_stage_model.quantize(x_recon)
+ model_mean, posterior_variance, posterior_log_variance = self.q_posterior(x_start=x_recon, x_t=x, t=t)
+ if return_codebook_ids:
+ return model_mean, posterior_variance, posterior_log_variance, logits
+ elif return_x0:
+ return model_mean, posterior_variance, posterior_log_variance, x_recon
+ else:
+ return model_mean, posterior_variance, posterior_log_variance
+
+ @torch.no_grad()
+ def p_sample(self, x, c, t, clip_denoised=False, repeat_noise=False,
+ return_codebook_ids=False, quantize_denoised=False, return_x0=False,
+ temperature=1., noise_dropout=0., score_corrector=None, corrector_kwargs=None):
+ b, *_, device = *x.shape, x.device
+ outputs = self.p_mean_variance(x=x, c=c, t=t, clip_denoised=clip_denoised,
+ return_codebook_ids=return_codebook_ids,
+ quantize_denoised=quantize_denoised,
+ return_x0=return_x0,
+ score_corrector=score_corrector, corrector_kwargs=corrector_kwargs)
+ if return_codebook_ids:
+ raise DeprecationWarning("Support dropped.")
+ model_mean, _, model_log_variance, logits = outputs
+ elif return_x0:
+ model_mean, _, model_log_variance, x0 = outputs
+ else:
+ model_mean, _, model_log_variance = outputs
+
+ noise = noise_like(x.shape, device, repeat_noise) * temperature
+ if noise_dropout > 0.:
+ noise = torch.nn.functional.dropout(noise, p=noise_dropout)
+ # no noise when t == 0
+ nonzero_mask = (1 - (t == 0).float()).reshape(b, *((1,) * (len(x.shape) - 1)))
+
+ if return_codebook_ids:
+ return model_mean + nonzero_mask * (0.5 * model_log_variance).exp() * noise, logits.argmax(dim=1)
+ if return_x0:
+ return model_mean + nonzero_mask * (0.5 * model_log_variance).exp() * noise, x0
+ else:
+ return model_mean + nonzero_mask * (0.5 * model_log_variance).exp() * noise
+
+ @torch.no_grad()
+ def progressive_denoising(self, cond, shape, verbose=False, callback=None, quantize_denoised=False,
+ img_callback=None, mask=None, x0=None, temperature=1., noise_dropout=0.,
+ score_corrector=None, corrector_kwargs=None, batch_size=None, x_T=None, start_T=None,
+ log_every_t=None):
+ if not log_every_t:
+ log_every_t = self.log_every_t
+ timesteps = self.num_timesteps
+ if batch_size is not None:
+ b = batch_size if batch_size is not None else shape[0]
+ shape = [batch_size] + list(shape)
+ else:
+ b = batch_size = shape[0]
+ if x_T is None:
+ img = torch.randn(shape, device=self.device)
+ else:
+ img = x_T
+ intermediates = []
+ if cond is not None:
+ if isinstance(cond, dict):
+ cond = {key: cond[key][:batch_size] if not isinstance(cond[key], list) else
+ list(map(lambda x: x[:batch_size], cond[key])) for key in cond}
+ else:
+ cond = [c[:batch_size] for c in cond] if isinstance(cond, list) else cond[:batch_size]
+
+ if start_T is not None:
+ timesteps = min(timesteps, start_T)
+ iterator = tqdm(reversed(range(0, timesteps)), desc='Progressive Generation',
+ total=timesteps) if verbose else reversed(
+ range(0, timesteps))
+ if type(temperature) == float:
+ temperature = [temperature] * timesteps
+
+ for i in iterator:
+ ts = torch.full((b,), i, device=self.device, dtype=torch.long)
+ if self.shorten_cond_schedule:
+ assert self.model.conditioning_key != 'hybrid'
+ tc = self.cond_ids[ts].to(cond.device)
+ cond = self.q_sample(x_start=cond, t=tc, noise=torch.randn_like(cond))
+
+ img, x0_partial = self.p_sample(img, cond, ts,
+ clip_denoised=self.clip_denoised,
+ quantize_denoised=quantize_denoised, return_x0=True,
+ temperature=temperature[i], noise_dropout=noise_dropout,
+ score_corrector=score_corrector, corrector_kwargs=corrector_kwargs)
+ if mask is not None:
+ assert x0 is not None
+ img_orig = self.q_sample(x0, ts)
+ img = img_orig * mask + (1. - mask) * img
+
+ if i % log_every_t == 0 or i == timesteps - 1:
+ intermediates.append(x0_partial)
+ if callback: callback(i)
+ if img_callback: img_callback(img, i)
+ return img, intermediates
+
+ @torch.no_grad()
+ def p_sample_loop(self, cond, shape, return_intermediates=False,
+ x_T=None, verbose=False, callback=None, timesteps=None, quantize_denoised=False,
+ mask=None, x0=None, img_callback=None, start_T=None,
+ log_every_t=None):
+
+ if not log_every_t:
+ log_every_t = self.log_every_t
+ device = self.betas.device
+ b = shape[0]
+ if x_T is None:
+ img = torch.randn(shape, device=device)
+ else:
+ img = x_T
+
+ intermediates = [img]
+ if timesteps is None:
+ timesteps = self.num_timesteps
+
+ if start_T is not None:
+ timesteps = min(timesteps, start_T)
+ iterator = tqdm(reversed(range(0, timesteps)), desc='Sampling t', total=timesteps) if verbose else reversed(
+ range(0, timesteps))
+
+ if mask is not None:
+ assert x0 is not None
+ assert x0.shape[2:3] == mask.shape[2:3] # spatial size has to match
+
+ for i in iterator:
+ ts = torch.full((b,), i, device=device, dtype=torch.long)
+ if self.shorten_cond_schedule:
+ assert self.model.conditioning_key != 'hybrid'
+ tc = self.cond_ids[ts].to(cond.device)
+ cond = self.q_sample(x_start=cond, t=tc, noise=torch.randn_like(cond))
+
+ img = self.p_sample(img, cond, ts,
+ clip_denoised=self.clip_denoised,
+ quantize_denoised=quantize_denoised)
+ if mask is not None:
+ img_orig = self.q_sample(x0, ts)
+ img = img_orig * mask + (1. - mask) * img
+
+ if i % log_every_t == 0 or i == timesteps - 1:
+ intermediates.append(img)
+ if callback: callback(i)
+ if img_callback: img_callback(img, i)
+
+ if return_intermediates:
+ return img, intermediates
+ return img
+
+ @torch.no_grad()
+ def sample(self, cond, batch_size=16, return_intermediates=False, x_T=None,
+ verbose=False, timesteps=None, quantize_denoised=False,
+ mask=None, x0=None, shape=None, **kwargs):
+ if shape is None:
+ shape = (batch_size, self.channels, *self.image_size)
+ if cond is not None:
+ if isinstance(cond, dict):
+ cond = {key: cond[key][:batch_size] if not isinstance(cond[key], list) else
+ list(map(lambda x: x[:batch_size], cond[key])) for key in cond}
+ else:
+ cond = [c[:batch_size] for c in cond] if isinstance(cond, list) else cond[:batch_size]
+ return self.p_sample_loop(cond,
+ shape,
+ return_intermediates=return_intermediates, x_T=x_T,
+ verbose=verbose, timesteps=timesteps, quantize_denoised=quantize_denoised,
+ mask=mask, x0=x0)
+
+ @torch.no_grad()
+ def sample_log(self, cond, batch_size, ddim, ddim_steps, **kwargs):
+ if ddim:
+ ddim_sampler = DDIMSampler(self)
+ shape = (self.channels, *self.image_size)
+ samples, intermediates = ddim_sampler.sample(ddim_steps, batch_size,
+ shape, cond, verbose=self.verbose, **kwargs)
+
+ else:
+ samples, intermediates = self.sample(cond=cond, batch_size=batch_size,
+ return_intermediates=True, **kwargs)
+
+ return samples, intermediates
+
+ @torch.no_grad()
+ def log_images(self, batch, N=8, n_row=4, sample=True, ddim_steps=200, ddim_eta=1., return_keys=None,
+ quantize_denoised=False, inpaint=False, plot_denoise_rows=False, plot_progressive_rows=False,
+ plot_diffusion_rows=False, dset=None, **kwargs):
+
+ use_ddim = ddim_steps is not None
+
+ log = dict()
+ z, c, x, xrec, xc = self.get_input(batch, self.first_stage_key,
+ return_first_stage_outputs=True,
+ force_c_encode=True,
+ return_original_cond=True,
+ bs=N)
+
+ N = min(x.shape[0], N)
+ n_row = min(x.shape[0], n_row)
+ log["inputs"] = x
+ log["reconstruction"] = xrec
+ if self.model.conditioning_key is not None:
+ if hasattr(self.cond_stage_model, "decode"):
+ xc = self.cond_stage_model.decode(c)
+ log["conditioning"] = xc
+ elif self.cond_stage_key in ["caption"]:
+ xc = log_txt_as_img((x.shape[2], x.shape[3]), batch["caption"])
+ log["conditioning"] = xc
+ elif self.cond_stage_key in ['class_label']:
+ xc = log_txt_as_img((x.shape[2], x.shape[3]), batch["human_label"])
+ log['conditioning'] = xc
+ elif self.cond_stage_key in ['camera']:
+ if isinstance(batch["camera"], list):
+ xc = torch.cat(batch["camera"], -1)
+ else:
+ xc = batch["camera"].permute(0, 2, 3, 1, 4)
+ xc = xc.reshape(*xc.shape[:3], -1) * 2. - 1.
+ log['conditioning'] = xc
+ elif isimage(xc):
+ log["conditioning"] = xc
+ if ismap(xc):
+ if dset is None:
+ key = 'train' if self.training else 'validation'
+ dset = self.trainer.datamodule.datasets[key].data
+ label2rgb = torch.from_numpy(dset.label2rgb).to(self.device) / 127.5 - 1.
+ # log["original_conditioning"] = self.to_rgb(xc)
+ log["original_conditioning"] = label2rgb[xc.argmax(1)].permute(0, 3, 1, 2)
+
+ if plot_diffusion_rows:
+ # get diffusion row
+ diffusion_row = list()
+ z_start = z[:n_row]
+ for t in range(self.num_timesteps):
+ if t % self.log_every_t == 0 or t == self.num_timesteps - 1:
+ t = repeat(torch.tensor([t]), '1 -> b', b=n_row)
+ t = t.to(self.device).long()
+ noise = torch.randn_like(z_start)
+ z_noisy = self.q_sample(x_start=z_start, t=t, noise=noise)
+ diffusion_row.append(self.decode_first_stage(z_noisy))
+
+ diffusion_row = torch.stack(diffusion_row) # n_log_step, n_row, C, H, W
+ diffusion_grid = rearrange(diffusion_row, 'n b c h w -> b n c h w')
+ diffusion_grid = rearrange(diffusion_grid, 'b n c h w -> (b n) c h w')
+ diffusion_grid = make_grid(diffusion_grid, nrow=diffusion_row.shape[0])
+ log["diffusion_row"] = diffusion_grid
+
+ if sample:
+ # get denoise row
+ with self.ema_scope("Plotting"):
+ samples, z_denoise_row = self.sample_log(cond=c, batch_size=N, ddim=use_ddim,
+ ddim_steps=ddim_steps, eta=ddim_eta)
+ x_samples = self.decode_first_stage(samples)
+ log["samples"] = x_samples
+ if plot_denoise_rows:
+ denoise_grid = self._get_denoise_row_from_list(z_denoise_row)
+ log["denoise_row"] = denoise_grid
+
+ if quantize_denoised and not isinstance(self.first_stage_model, AutoencoderKL) and not isinstance(
+ self.first_stage_model, IdentityFirstStage):
+ # also display when quantizing x0 while sampling
+ with self.ema_scope("Plotting Quantized Denoised"):
+ samples, z_denoise_row = self.sample_log(cond=c, batch_size=N, ddim=use_ddim,
+ ddim_steps=ddim_steps, eta=ddim_eta,
+ quantize_denoised=True)
+ x_samples = self.decode_first_stage(samples.to(self.device))
+ log["samples_x0_quantized"] = x_samples
+
+ if inpaint:
+ # make a simple center square
+ b, h, w = z.shape[0], z.shape[2], z.shape[3]
+ mask = torch.ones(N, h, w).to(self.device)
+ # zeros will be filled in
+ mask[:, h // 4:3 * h // 4, w // 4:3 * w // 4] = 0.
+ mask = mask[:, None, ...]
+ with self.ema_scope("Plotting Inpaint"):
+ samples, _ = self.sample_log(cond=c, batch_size=N, ddim=use_ddim, eta=ddim_eta,
+ ddim_steps=ddim_steps, x0=z[:N], mask=mask)
+ x_samples = self.decode_first_stage(samples.to(self.device))
+ log["samples_inpainting"] = x_samples
+ log["mask"] = mask
+
+ # outpaint
+ with self.ema_scope("Plotting Outpaint"):
+ samples, _ = self.sample_log(cond=c, batch_size=N, ddim=use_ddim, eta=ddim_eta,
+ ddim_steps=ddim_steps, x0=z[:N], mask=mask)
+ x_samples = self.decode_first_stage(samples.to(self.device))
+ log["samples_outpainting"] = x_samples
+
+ if plot_progressive_rows:
+ with self.ema_scope("Plotting Progressives"):
+ img, progressives = self.progressive_denoising(c, shape=(self.channels, *self.image_size), batch_size=N)
+ prog_row = self._get_denoise_row_from_list(progressives, desc="Progressive Generation")
+ log["progressive_row"] = prog_row
+
+ if return_keys:
+ if np.intersect1d(list(log.keys()), return_keys).shape[0] == 0:
+ return log
+ else:
+ return {key: log[key] for key in return_keys}
+ return log
+
+ def configure_optimizers(self):
+ lr = self.learning_rate
+ params = list(self.model.parameters())
+ if self.cond_stage_trainable:
+ self.print_fn(f"{self.__class__.__name__}: Also optimizing conditioner params!")
+ params = params + list(self.cond_stage_model.parameters())
+ if self.learn_logvar:
+ self.print_fn('Diffusion model optimizing logvar')
+ params.append(self.logvar)
+ opt = torch.optim.AdamW(params, lr=lr)
+ if self.use_scheduler:
+ assert 'target' in self.scheduler_config
+ scheduler = instantiate_from_config(self.scheduler_config)
+
+ self.print_fn("Setting up LambdaLR scheduler...")
+ scheduler = [
+ {
+ 'scheduler': LambdaLR(opt, lr_lambda=scheduler.schedule),
+ 'interval': 'step',
+ 'frequency': 1
+ }]
+ return [opt], scheduler
+ return opt
+
+ @torch.no_grad()
+ def to_rgb(self, x):
+ x = x.float()
+ if not hasattr(self, "colorize"):
+ self.colorize = torch.randn(3, x.shape[1], 1, 1).to(x)
+ x = nn.functional.conv2d(x, weight=self.colorize)
+ x = 2. * (x - x.min()) / (x.max() - x.min()) - 1.
+ return x
+
+
+class DiffusionWrapper(pl.LightningModule):
+ def __init__(self, diff_model_config, conditioning_key):
+ super().__init__()
+ self.diffusion_model = instantiate_from_config(diff_model_config)
+ self.conditioning_key = conditioning_key
+ assert self.conditioning_key in [None, 'concat', 'crossattn', 'hybrid', 'adm']
+
+ def forward(self, x, t, c_concat: list = None, c_crossattn: list = None):
+ if self.conditioning_key is None:
+ out = self.diffusion_model(x, t)
+ elif self.conditioning_key == 'concat':
+ xc = torch.cat([x] + c_concat, dim=1)
+ out = self.diffusion_model(xc, t)
+ elif self.conditioning_key == 'crossattn':
+ cc = torch.cat(c_crossattn, 1)
+ out = self.diffusion_model(x, t, context=cc)
+ elif self.conditioning_key == 'hybrid':
+ xc = torch.cat([x] + c_concat, dim=1)
+ cc = torch.cat(c_crossattn, 1)
+ out = self.diffusion_model(xc, t, context=cc)
+ elif self.conditioning_key == 'adm':
+ cc = c_crossattn[0]
+ out = self.diffusion_model(x, t, y=cc)
+ else:
+ raise NotImplementedError()
+
+ return out
+
+
+class Layout2ImgDiffusion(LatentDiffusion):
+ def __init__(self, cond_stage_key, *args, **kwargs):
+ assert cond_stage_key in ['bbox', 'center'], 'Layout2ImgDiffusion only for cond_stage_key="bbox" or "center"'
+ super().__init__(cond_stage_key=cond_stage_key, *args, **kwargs)
+
+ def log_images(self, batch, N=8, dset=None, *args, **kwargs):
+ logs = super().log_images(batch=batch, N=N, *args, **kwargs)
+
+ key = 'train' if self.training else 'validation'
+ if dset is None:
+ dset = self.trainer.datamodule.datasets[key].data
+ mapper = dset.conditional_builders[self.cond_stage_key]
+ H, W = batch['image'].shape[2:]
+
+ bbox_imgs = []
+ map_fn = lambda catno: dset.get_textual_label_for_category_id(catno)
+ for tknzd_bbox in batch[self.cond_stage_key][:N]:
+ bboximg = mapper.plot(tknzd_bbox.detach().cpu(), map_fn, (W, H))
+ bbox_imgs.append(bboximg)
+
+ cond_img = torch.stack(bbox_imgs, dim=0)
+ logs['bbox_image'] = cond_img
+ return logs
diff --git a/lidm/models/diffusion/plms.py b/lidm/models/diffusion/plms.py
new file mode 100644
index 0000000000000000000000000000000000000000..b68ab413ca361d8bc8ba815c12ad2dbb2144d728
--- /dev/null
+++ b/lidm/models/diffusion/plms.py
@@ -0,0 +1,236 @@
+"""SAMPLING ONLY."""
+
+import torch
+import numpy as np
+from tqdm import tqdm
+from functools import partial
+
+from ...modules.diffusion.util import make_ddim_sampling_parameters, make_ddim_timesteps, noise_like
+
+
+class PLMSSampler(object):
+ def __init__(self, model, schedule="linear", **kwargs):
+ super().__init__()
+ self.model = model
+ self.ddpm_num_timesteps = model.num_timesteps
+ self.schedule = schedule
+
+ def register_buffer(self, name, attr):
+ if type(attr) == torch.Tensor:
+ if attr.device != torch.device("cuda"):
+ attr = attr.to(torch.device("cuda"))
+ setattr(self, name, attr)
+
+ def make_schedule(self, ddim_num_steps, ddim_discretize="uniform", ddim_eta=0., verbose=False):
+ if ddim_eta != 0:
+ raise ValueError('ddim_eta must be 0 for PLMS')
+ self.ddim_timesteps = make_ddim_timesteps(ddim_discr_method=ddim_discretize, num_ddim_timesteps=ddim_num_steps,
+ num_ddpm_timesteps=self.ddpm_num_timesteps,verbose=verbose)
+ alphas_cumprod = self.model.alphas_cumprod
+ assert alphas_cumprod.shape[0] == self.ddpm_num_timesteps, 'alphas have to be defined for each timestep'
+ to_torch = lambda x: x.clone().detach().to(torch.float32).to(self.model.device)
+
+ self.register_buffer('betas', to_torch(self.model.betas))
+ self.register_buffer('alphas_cumprod', to_torch(alphas_cumprod))
+ self.register_buffer('alphas_cumprod_prev', to_torch(self.model.alphas_cumprod_prev))
+
+ # calculations for diffusion q(x_t | x_{t-1}) and others
+ self.register_buffer('sqrt_alphas_cumprod', to_torch(np.sqrt(alphas_cumprod.cpu())))
+ self.register_buffer('sqrt_one_minus_alphas_cumprod', to_torch(np.sqrt(1. - alphas_cumprod.cpu())))
+ self.register_buffer('log_one_minus_alphas_cumprod', to_torch(np.log(1. - alphas_cumprod.cpu())))
+ self.register_buffer('sqrt_recip_alphas_cumprod', to_torch(np.sqrt(1. / alphas_cumprod.cpu())))
+ self.register_buffer('sqrt_recipm1_alphas_cumprod', to_torch(np.sqrt(1. / alphas_cumprod.cpu() - 1)))
+
+ # ddim sampling parameters
+ ddim_sigmas, ddim_alphas, ddim_alphas_prev = make_ddim_sampling_parameters(alphacums=alphas_cumprod.cpu(),
+ ddim_timesteps=self.ddim_timesteps,
+ eta=ddim_eta,verbose=verbose)
+ self.register_buffer('ddim_sigmas', ddim_sigmas)
+ self.register_buffer('ddim_alphas', ddim_alphas)
+ self.register_buffer('ddim_alphas_prev', ddim_alphas_prev)
+ self.register_buffer('ddim_sqrt_one_minus_alphas', np.sqrt(1. - ddim_alphas))
+ sigmas_for_original_sampling_steps = ddim_eta * torch.sqrt(
+ (1 - self.alphas_cumprod_prev) / (1 - self.alphas_cumprod) * (
+ 1 - self.alphas_cumprod / self.alphas_cumprod_prev))
+ self.register_buffer('ddim_sigmas_for_original_num_steps', sigmas_for_original_sampling_steps)
+
+ @torch.no_grad()
+ def sample(self,
+ S,
+ batch_size,
+ shape,
+ conditioning=None,
+ callback=None,
+ normals_sequence=None,
+ img_callback=None,
+ quantize_x0=False,
+ eta=0.,
+ mask=None,
+ x0=None,
+ temperature=1.,
+ noise_dropout=0.,
+ score_corrector=None,
+ corrector_kwargs=None,
+ verbose=False,
+ x_T=None,
+ log_every_t=100,
+ unconditional_guidance_scale=1.,
+ unconditional_conditioning=None,
+ # this has to come in the same format as the conditioning, # e.g. as encoded tokens, ...
+ **kwargs
+ ):
+ if conditioning is not None:
+ if isinstance(conditioning, dict):
+ cbs = conditioning[list(conditioning.keys())[0]].shape[0]
+ if cbs != batch_size:
+ print(f"Warning: Got {cbs} conditionings but batch-size is {batch_size}")
+ else:
+ if conditioning.shape[0] != batch_size:
+ print(f"Warning: Got {conditioning.shape[0]} conditionings but batch-size is {batch_size}")
+
+ self.make_schedule(ddim_num_steps=S, ddim_eta=eta, verbose=verbose)
+ # sampling
+ C, H, W = shape
+ size = (batch_size, C, H, W)
+ print(f'Data shape for PLMS sampling is {size}')
+
+ samples, intermediates = self.plms_sampling(conditioning, size,
+ callback=callback,
+ img_callback=img_callback,
+ quantize_denoised=quantize_x0,
+ mask=mask, x0=x0,
+ ddim_use_original_steps=False,
+ noise_dropout=noise_dropout,
+ temperature=temperature,
+ score_corrector=score_corrector,
+ corrector_kwargs=corrector_kwargs,
+ x_T=x_T,
+ log_every_t=log_every_t,
+ unconditional_guidance_scale=unconditional_guidance_scale,
+ unconditional_conditioning=unconditional_conditioning,
+ )
+ return samples, intermediates
+
+ @torch.no_grad()
+ def plms_sampling(self, cond, shape,
+ x_T=None, ddim_use_original_steps=False,
+ callback=None, timesteps=None, quantize_denoised=False,
+ mask=None, x0=None, img_callback=None, log_every_t=100,
+ temperature=1., noise_dropout=0., score_corrector=None, corrector_kwargs=None,
+ unconditional_guidance_scale=1., unconditional_conditioning=None,):
+ device = self.model.betas.device
+ b = shape[0]
+ if x_T is None:
+ img = torch.randn(shape, device=device)
+ else:
+ img = x_T
+
+ if timesteps is None:
+ timesteps = self.ddpm_num_timesteps if ddim_use_original_steps else self.ddim_timesteps
+ elif timesteps is not None and not ddim_use_original_steps:
+ subset_end = int(min(timesteps / self.ddim_timesteps.shape[0], 1) * self.ddim_timesteps.shape[0]) - 1
+ timesteps = self.ddim_timesteps[:subset_end]
+
+ intermediates = {'x_inter': [img], 'pred_x0': [img]}
+ time_range = list(reversed(range(0,timesteps))) if ddim_use_original_steps else np.flip(timesteps)
+ total_steps = timesteps if ddim_use_original_steps else timesteps.shape[0]
+ print(f"Running PLMS Sampling with {total_steps} timesteps")
+
+ iterator = tqdm(time_range, desc='PLMS Sampler', total=total_steps)
+ old_eps = []
+
+ for i, step in enumerate(iterator):
+ index = total_steps - i - 1
+ ts = torch.full((b,), step, device=device, dtype=torch.long)
+ ts_next = torch.full((b,), time_range[min(i + 1, len(time_range) - 1)], device=device, dtype=torch.long)
+
+ if mask is not None:
+ assert x0 is not None
+ img_orig = self.model.q_sample(x0, ts) # TODO: deterministic forward pass?
+ img = img_orig * mask + (1. - mask) * img
+
+ outs = self.p_sample_plms(img, cond, ts, index=index, use_original_steps=ddim_use_original_steps,
+ quantize_denoised=quantize_denoised, temperature=temperature,
+ noise_dropout=noise_dropout, score_corrector=score_corrector,
+ corrector_kwargs=corrector_kwargs,
+ unconditional_guidance_scale=unconditional_guidance_scale,
+ unconditional_conditioning=unconditional_conditioning,
+ old_eps=old_eps, t_next=ts_next)
+ img, pred_x0, e_t = outs
+ old_eps.append(e_t)
+ if len(old_eps) >= 4:
+ old_eps.pop(0)
+ if callback: callback(i)
+ if img_callback: img_callback(pred_x0, i)
+
+ if index % log_every_t == 0 or index == total_steps - 1:
+ intermediates['x_inter'].append(img)
+ intermediates['pred_x0'].append(pred_x0)
+
+ return img, intermediates
+
+ @torch.no_grad()
+ def p_sample_plms(self, x, c, t, index, repeat_noise=False, use_original_steps=False, quantize_denoised=False,
+ temperature=1., noise_dropout=0., score_corrector=None, corrector_kwargs=None,
+ unconditional_guidance_scale=1., unconditional_conditioning=None, old_eps=None, t_next=None):
+ b, *_, device = *x.shape, x.device
+
+ def get_model_output(x, t):
+ if unconditional_conditioning is None or unconditional_guidance_scale == 1.:
+ e_t = self.model.apply_model(x, t, c)
+ else:
+ x_in = torch.cat([x] * 2)
+ t_in = torch.cat([t] * 2)
+ c_in = torch.cat([unconditional_conditioning, c])
+ e_t_uncond, e_t = self.model.apply_model(x_in, t_in, c_in).chunk(2)
+ e_t = e_t_uncond + unconditional_guidance_scale * (e_t - e_t_uncond)
+
+ if score_corrector is not None:
+ assert self.model.parameterization == "eps"
+ e_t = score_corrector.modify_score(self.model, e_t, x, t, c, **corrector_kwargs)
+
+ return e_t
+
+ alphas = self.model.alphas_cumprod if use_original_steps else self.ddim_alphas
+ alphas_prev = self.model.alphas_cumprod_prev if use_original_steps else self.ddim_alphas_prev
+ sqrt_one_minus_alphas = self.model.sqrt_one_minus_alphas_cumprod if use_original_steps else self.ddim_sqrt_one_minus_alphas
+ sigmas = self.model.ddim_sigmas_for_original_num_steps if use_original_steps else self.ddim_sigmas
+
+ def get_x_prev_and_pred_x0(e_t, index):
+ # select parameters corresponding to the currently considered timestep
+ a_t = torch.full((b, 1, 1, 1), alphas[index], device=device)
+ a_prev = torch.full((b, 1, 1, 1), alphas_prev[index], device=device)
+ sigma_t = torch.full((b, 1, 1, 1), sigmas[index], device=device)
+ sqrt_one_minus_at = torch.full((b, 1, 1, 1), sqrt_one_minus_alphas[index],device=device)
+
+ # current prediction for x_0
+ pred_x0 = (x - sqrt_one_minus_at * e_t) / a_t.sqrt()
+ if quantize_denoised:
+ pred_x0, _, *_ = self.model.first_stage_model.quantize(pred_x0)
+ # direction pointing to x_t
+ dir_xt = (1. - a_prev - sigma_t**2).sqrt() * e_t
+ noise = sigma_t * noise_like(x.shape, device, repeat_noise) * temperature
+ if noise_dropout > 0.:
+ noise = torch.nn.functional.dropout(noise, p=noise_dropout)
+ x_prev = a_prev.sqrt() * pred_x0 + dir_xt + noise
+ return x_prev, pred_x0
+
+ e_t = get_model_output(x, t)
+ if len(old_eps) == 0:
+ # Pseudo Improved Euler (2nd order)
+ x_prev, pred_x0 = get_x_prev_and_pred_x0(e_t, index)
+ e_t_next = get_model_output(x_prev, t_next)
+ e_t_prime = (e_t + e_t_next) / 2
+ elif len(old_eps) == 1:
+ # 2nd order Pseudo Linear Multistep (Adams-Bashforth)
+ e_t_prime = (3 * e_t - old_eps[-1]) / 2
+ elif len(old_eps) == 2:
+ # 3nd order Pseudo Linear Multistep (Adams-Bashforth)
+ e_t_prime = (23 * e_t - 16 * old_eps[-1] + 5 * old_eps[-2]) / 12
+ elif len(old_eps) >= 3:
+ # 4nd order Pseudo Linear Multistep (Adams-Bashforth)
+ e_t_prime = (55 * e_t - 59 * old_eps[-1] + 37 * old_eps[-2] - 9 * old_eps[-3]) / 24
+
+ x_prev, pred_x0 = get_x_prev_and_pred_x0(e_t_prime, index)
+
+ return x_prev, pred_x0, e_t
diff --git a/lidm/modules/__init__.py b/lidm/modules/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/lidm/modules/attention.py b/lidm/modules/attention.py
new file mode 100644
index 0000000000000000000000000000000000000000..83708247da4a5d293aa6bbc83a15783eb32282df
--- /dev/null
+++ b/lidm/modules/attention.py
@@ -0,0 +1,261 @@
+from inspect import isfunction
+import math
+import torch
+import torch.nn.functional as F
+from torch import nn, einsum
+from einops import rearrange, repeat
+
+from .basic import checkpoint
+
+
+def exists(val):
+ return val is not None
+
+
+def uniq(arr):
+ return{el: True for el in arr}.keys()
+
+
+def default(val, d):
+ if exists(val):
+ return val
+ return d() if isfunction(d) else d
+
+
+def max_neg_value(t):
+ return -torch.finfo(t.dtype).max
+
+
+def init_(tensor):
+ dim = tensor.shape[-1]
+ std = 1 / math.sqrt(dim)
+ tensor.uniform_(-std, std)
+ return tensor
+
+
+# feedforward
+class GEGLU(nn.Module):
+ def __init__(self, dim_in, dim_out):
+ super().__init__()
+ self.proj = nn.Linear(dim_in, dim_out * 2)
+
+ def forward(self, x):
+ x, gate = self.proj(x).chunk(2, dim=-1)
+ return x * F.gelu(gate)
+
+
+class FeedForward(nn.Module):
+ def __init__(self, dim, dim_out=None, mult=4, glu=False, dropout=0.):
+ super().__init__()
+ inner_dim = int(dim * mult)
+ dim_out = default(dim_out, dim)
+ project_in = nn.Sequential(
+ nn.Linear(dim, inner_dim),
+ nn.GELU()
+ ) if not glu else GEGLU(dim, inner_dim)
+
+ self.net = nn.Sequential(
+ project_in,
+ nn.Dropout(dropout),
+ nn.Linear(inner_dim, dim_out)
+ )
+
+ def forward(self, x):
+ return self.net(x)
+
+
+def zero_module(module):
+ """
+ Zero out the parameters of a module and return it.
+ """
+ for p in module.parameters():
+ p.detach().zero_()
+ return module
+
+
+def Normalize(in_channels):
+ return torch.nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True)
+
+
+class LinearAttention(nn.Module):
+ def __init__(self, dim, heads=4, dim_head=32):
+ super().__init__()
+ self.heads = heads
+ hidden_dim = dim_head * heads
+ self.to_qkv = nn.Conv2d(dim, hidden_dim * 3, 1, bias=False)
+ self.to_out = nn.Conv2d(hidden_dim, dim, 1)
+
+ def forward(self, x):
+ b, c, h, w = x.shape
+ qkv = self.to_qkv(x)
+ q, k, v = rearrange(qkv, 'b (qkv heads c) h w -> qkv b heads c (h w)', heads = self.heads, qkv=3)
+ k = k.softmax(dim=-1)
+ context = torch.einsum('bhdn,bhen->bhde', k, v)
+ out = torch.einsum('bhde,bhdn->bhen', context, q)
+ out = rearrange(out, 'b heads c (h w) -> b (heads c) h w', heads=self.heads, h=h, w=w)
+ return self.to_out(out)
+
+
+class SpatialSelfAttention(nn.Module):
+ def __init__(self, in_channels):
+ super().__init__()
+ self.in_channels = in_channels
+
+ self.norm = Normalize(in_channels)
+ self.q = torch.nn.Conv2d(in_channels,
+ in_channels,
+ kernel_size=1,
+ stride=1,
+ padding=0)
+ self.k = torch.nn.Conv2d(in_channels,
+ in_channels,
+ kernel_size=1,
+ stride=1,
+ padding=0)
+ self.v = torch.nn.Conv2d(in_channels,
+ in_channels,
+ kernel_size=1,
+ stride=1,
+ padding=0)
+ self.proj_out = torch.nn.Conv2d(in_channels,
+ in_channels,
+ kernel_size=1,
+ stride=1,
+ padding=0)
+
+ def forward(self, x):
+ h_ = x
+ h_ = self.norm(h_)
+ q = self.q(h_)
+ k = self.k(h_)
+ v = self.v(h_)
+
+ # compute attention
+ b,c,h,w = q.shape
+ q = rearrange(q, 'b c h w -> b (h w) c')
+ k = rearrange(k, 'b c h w -> b c (h w)')
+ w_ = torch.einsum('bij,bjk->bik', q, k)
+
+ w_ = w_ * (int(c)**(-0.5))
+ w_ = torch.nn.functional.softmax(w_, dim=2)
+
+ # attend to values
+ v = rearrange(v, 'b c h w -> b c (h w)')
+ w_ = rearrange(w_, 'b i j -> b j i')
+ h_ = torch.einsum('bij,bjk->bik', v, w_)
+ h_ = rearrange(h_, 'b c (h w) -> b c h w', h=h)
+ h_ = self.proj_out(h_)
+
+ return x+h_
+
+
+class CrossAttention(nn.Module):
+ def __init__(self, query_dim, context_dim=None, heads=8, dim_head=64, dropout=0.):
+ super().__init__()
+ inner_dim = dim_head * heads
+ context_dim = default(context_dim, query_dim)
+
+ self.scale = dim_head ** -0.5
+ self.heads = heads
+
+ self.to_q = nn.Linear(query_dim, inner_dim, bias=False)
+ self.to_k = nn.Linear(context_dim, inner_dim, bias=False)
+ self.to_v = nn.Linear(context_dim, inner_dim, bias=False)
+
+ self.to_out = nn.Sequential(
+ nn.Linear(inner_dim, query_dim),
+ nn.Dropout(dropout)
+ )
+
+ def forward(self, x, context=None, mask=None):
+ h = self.heads
+
+ q = self.to_q(x)
+ context = default(context, x)
+ k = self.to_k(context)
+ v = self.to_v(context)
+
+ q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> (b h) n d', h=h), (q, k, v))
+
+ sim = einsum('b i d, b j d -> b i j', q, k) * self.scale
+
+ if exists(mask):
+ mask = rearrange(mask, 'b ... -> b (...)')
+ max_neg_value = -torch.finfo(sim.dtype).max
+ mask = repeat(mask, 'b j -> (b h) () j', h=h)
+ sim.masked_fill_(~mask, max_neg_value)
+
+ # attention, what we cannot get enough of
+ attn = sim.softmax(dim=-1)
+
+ out = einsum('b i j, b j d -> b i d', attn, v)
+ out = rearrange(out, '(b h) n d -> b n (h d)', h=h)
+ return self.to_out(out)
+
+
+class BasicTransformerBlock(nn.Module):
+ def __init__(self, dim, n_heads, d_head, dropout=0., context_dim=None, gated_ff=True, checkpoint=True):
+ super().__init__()
+ self.attn1 = CrossAttention(query_dim=dim, heads=n_heads, dim_head=d_head, dropout=dropout) # is a self-attention
+ self.ff = FeedForward(dim, dropout=dropout, glu=gated_ff)
+ self.attn2 = CrossAttention(query_dim=dim, context_dim=context_dim,
+ heads=n_heads, dim_head=d_head, dropout=dropout) # is self-attn if context is none
+ self.norm1 = nn.LayerNorm(dim)
+ self.norm2 = nn.LayerNorm(dim)
+ self.norm3 = nn.LayerNorm(dim)
+ self.checkpoint = checkpoint
+
+ def forward(self, x, context=None):
+ return checkpoint(self._forward, (x, context), self.parameters(), self.checkpoint)
+
+ def _forward(self, x, context=None):
+ x = self.attn1(self.norm1(x)) + x
+ x = self.attn2(self.norm2(x), context=context) + x
+ x = self.ff(self.norm3(x)) + x
+ return x
+
+
+class SpatialTransformer(nn.Module):
+ """
+ Transformer block for image-like data.
+ First, project the input (aka embedding)
+ and reshape to b, t, d.
+ Then apply standard transformer action.
+ Finally, reshape to image
+ """
+ def __init__(self, in_channels, n_heads, d_head,
+ depth=1, dropout=0., context_dim=None):
+ super().__init__()
+ self.in_channels = in_channels
+ inner_dim = n_heads * d_head
+ self.norm = Normalize(in_channels)
+
+ self.proj_in = nn.Conv2d(in_channels,
+ inner_dim,
+ kernel_size=1,
+ stride=1,
+ padding=0)
+
+ self.transformer_blocks = nn.ModuleList(
+ [BasicTransformerBlock(inner_dim, n_heads, d_head, dropout=dropout, context_dim=context_dim)
+ for d in range(depth)]
+ )
+
+ self.proj_out = zero_module(nn.Conv2d(inner_dim,
+ in_channels,
+ kernel_size=1,
+ stride=1,
+ padding=0))
+
+ def forward(self, x, context=None):
+ # note: if no context is given, cross-attention defaults to self-attention
+ b, c, h, w = x.shape
+ x_in = x
+ x = self.norm(x)
+ x = self.proj_in(x)
+ x = rearrange(x, 'b c h w -> b (h w) c')
+ for block in self.transformer_blocks:
+ x = block(x, context=context)
+ x = rearrange(x, 'b (h w) c -> b c h w', h=h, w=w)
+ x = self.proj_out(x)
+ return x + x_in
\ No newline at end of file
diff --git a/lidm/modules/basic.py b/lidm/modules/basic.py
new file mode 100644
index 0000000000000000000000000000000000000000..e0571d7a5f33ef7c26b219e20d8f31859936b5ed
--- /dev/null
+++ b/lidm/modules/basic.py
@@ -0,0 +1,392 @@
+# adopted from
+# https://github.com/openai/improved-diffusion/blob/main/improved_diffusion/gaussian_diffusion.py
+# and
+# https://github.com/lucidrains/denoising-diffusion-pytorch/blob/7706bdfc6f527f58d33f84b7b522e61e6e3164b3/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py
+# and
+# https://github.com/openai/guided-diffusion/blob/0ba878e517b276c45d1195eb29f6f5f72659a05b/guided_diffusion/nn.py
+#
+# thanks!
+
+
+import math
+import torch
+import torch.nn as nn
+import numpy as np
+from einops import repeat
+from torch import Tensor
+
+from ..utils.misc_utils import instantiate_from_config, print_fn
+
+
+class CircularPad(nn.Module):
+ def __init__(self, pad_size):
+ super(CircularPad, self).__init__()
+ h1, h2, v1, v2 = pad_size
+ self.h_pad, self.v_pad = (h1, h2, 0, 0), (0, 0, v1, v2)
+
+ def forward(self, x):
+ if sum(self.h_pad) > 0:
+ x = nn.functional.pad(x, self.h_pad, mode="circular") # horizontal pad
+ if sum(self.v_pad) > 0:
+ x = nn.functional.pad(x, self.v_pad, mode="constant") # vertical pad
+ return x
+
+
+class CircularConv2d(nn.Conv2d):
+ def __init__(self, *args, **kwargs):
+ if 'padding' in kwargs:
+ self.is_pad = True
+ if isinstance(kwargs['padding'], int):
+ h1 = h2 = v1 = v2 = kwargs['padding']
+ elif isinstance(kwargs['padding'], tuple):
+ h1, h2, v1, v2 = kwargs['padding']
+ else:
+ raise NotImplementedError
+ self.h_pad, self.v_pad = (h1, h2, 0, 0), (0, 0, v1, v2)
+ del kwargs['padding']
+ else:
+ self.is_pad = False
+
+ super().__init__(*args, **kwargs)
+
+ def forward(self, x: Tensor) -> Tensor:
+ if self.is_pad:
+ if sum(self.h_pad) > 0:
+ x = nn.functional.pad(x, self.h_pad, mode="circular") # horizontal pad
+ if sum(self.v_pad) > 0:
+ x = nn.functional.pad(x, self.v_pad, mode="constant") # vertical pad
+ x = self._conv_forward(x, self.weight, self.bias)
+ return x
+
+
+class ActNorm(nn.Module):
+ def __init__(self, num_features, logdet=False, affine=True,
+ allow_reverse_init=False):
+ assert affine
+ super().__init__()
+ self.logdet = logdet
+ self.loc = nn.Parameter(torch.zeros(1, num_features, 1, 1))
+ self.scale = nn.Parameter(torch.ones(1, num_features, 1, 1))
+ self.allow_reverse_init = allow_reverse_init
+
+ self.register_buffer('initialized', torch.tensor(0, dtype=torch.uint8))
+
+ def initialize(self, input):
+ with torch.no_grad():
+ flatten = input.permute(1, 0, 2, 3).contiguous().view(input.shape[1], -1)
+ mean = (
+ flatten.mean(1)
+ .unsqueeze(1)
+ .unsqueeze(2)
+ .unsqueeze(3)
+ .permute(1, 0, 2, 3)
+ )
+ std = (
+ flatten.std(1)
+ .unsqueeze(1)
+ .unsqueeze(2)
+ .unsqueeze(3)
+ .permute(1, 0, 2, 3)
+ )
+
+ self.loc.data.copy_(-mean)
+ self.scale.data.copy_(1 / (std + 1e-6))
+
+ def forward(self, input, reverse=False):
+ if reverse:
+ return self.reverse(input)
+ if len(input.shape) == 2:
+ input = input[:,:,None,None]
+ squeeze = True
+ else:
+ squeeze = False
+
+ _, _, height, width = input.shape
+
+ if self.training and self.initialized.item() == 0:
+ self.initialize(input)
+ self.initialized.fill_(1)
+
+ h = self.scale * (input + self.loc)
+
+ if squeeze:
+ h = h.squeeze(-1).squeeze(-1)
+
+ if self.logdet:
+ log_abs = torch.log(torch.abs(self.scale))
+ logdet = height*width*torch.sum(log_abs)
+ logdet = logdet * torch.ones(input.shape[0]).to(input)
+ return h, logdet
+
+ return h
+
+ def reverse(self, output):
+ if self.training and self.initialized.item() == 0:
+ if not self.allow_reverse_init:
+ raise RuntimeError(
+ "Initializing ActNorm in reverse direction is "
+ "disabled by default. Use allow_reverse_init=True to enable."
+ )
+ else:
+ self.initialize(output)
+ self.initialized.fill_(1)
+
+ if len(output.shape) == 2:
+ output = output[:, :, None, None]
+ squeeze = True
+ else:
+ squeeze = False
+
+ h = output / self.scale - self.loc
+
+ if squeeze:
+ h = h.squeeze(-1).squeeze(-1)
+ return h
+
+
+def make_beta_schedule(schedule, n_timestep, linear_start=1e-4, linear_end=2e-2, cosine_s=8e-3):
+ if schedule == "linear":
+ betas = (
+ torch.linspace(linear_start ** 0.5, linear_end ** 0.5, n_timestep, dtype=torch.float64) ** 2
+ )
+
+ elif schedule == "cosine":
+ timesteps = (
+ torch.arange(n_timestep + 1, dtype=torch.float64) / n_timestep + cosine_s
+ )
+ alphas = timesteps / (1 + cosine_s) * np.pi / 2
+ alphas = torch.cos(alphas).pow(2)
+ alphas = alphas / alphas[0]
+ betas = 1 - alphas[1:] / alphas[:-1]
+ betas = np.clip(betas, a_min=0, a_max=0.999)
+
+ elif schedule == "sqrt_linear":
+ betas = torch.linspace(linear_start, linear_end, n_timestep, dtype=torch.float64)
+ elif schedule == "sqrt":
+ betas = torch.linspace(linear_start, linear_end, n_timestep, dtype=torch.float64) ** 0.5
+ else:
+ raise ValueError(f"schedule '{schedule}' unknown.")
+ return betas.numpy()
+
+
+def make_ddim_timesteps(ddim_discr_method, num_ddim_timesteps, num_ddpm_timesteps, verbose=False):
+ if ddim_discr_method == 'uniform':
+ c = num_ddpm_timesteps // num_ddim_timesteps
+ ddim_timesteps = np.asarray(list(range(0, num_ddpm_timesteps, c)))
+ elif ddim_discr_method == 'quad':
+ ddim_timesteps = ((np.linspace(0, np.sqrt(num_ddpm_timesteps * .8), num_ddim_timesteps)) ** 2).astype(int)
+ else:
+ raise NotImplementedError(f'There is no ddim discretization method called "{ddim_discr_method}"')
+
+ # assert ddim_timesteps.shape[0] == num_ddim_timesteps
+ # add one to get the final alpha values right (the ones from first scale to data during sampling)
+ steps_out = ddim_timesteps + 1
+ print_fn(f'Selected timesteps for ddim sampler: {steps_out}', False)
+ return steps_out
+
+
+def make_ddim_sampling_parameters(alphacums, ddim_timesteps, eta, verbose=False):
+ # select alphas for computing the variance schedule
+ alphas = alphacums[ddim_timesteps]
+ alphas_prev = np.asarray([alphacums[0]] + alphacums[ddim_timesteps[:-1]].tolist())
+
+ # according the the formula provided in https://arxiv.org/abs/2010.02502
+ sigmas = eta * np.sqrt((1 - alphas_prev) / (1 - alphas) * (1 - alphas / alphas_prev))
+ print_fn(f'Selected alphas for ddim sampler: a_t: {alphas}; a_(t-1): {alphas_prev}', False)
+ print_fn(f'For the chosen value of eta, which is {eta}, this results in the following sigma_t schedule for ddim sampler {sigmas}', False)
+ return sigmas, alphas, alphas_prev
+
+
+def betas_for_alpha_bar(num_diffusion_timesteps, alpha_bar, max_beta=0.999):
+ """
+ Create a beta schedule that discretizes the given alpha_t_bar function,
+ which defines the cumulative product of (1-beta) over time from t = [0,1].
+ :param num_diffusion_timesteps: the number of betas to produce.
+ :param alpha_bar: a lambda that takes an argument t from 0 to 1 and
+ produces the cumulative product of (1-beta) up to that
+ part of the diffusion process.
+ :param max_beta: the maximum beta to use; use values lower than 1 to
+ prevent singularities.
+ """
+ betas = []
+ for i in range(num_diffusion_timesteps):
+ t1 = i / num_diffusion_timesteps
+ t2 = (i + 1) / num_diffusion_timesteps
+ betas.append(min(1 - alpha_bar(t2) / alpha_bar(t1), max_beta))
+ return np.array(betas)
+
+
+def extract_into_tensor(a, t, x_shape):
+ b, *_ = t.shape
+ out = a.gather(-1, t)
+ return out.reshape(b, *((1,) * (len(x_shape) - 1)))
+
+
+def checkpoint(func, inputs, params, flag):
+ """
+ Evaluate a function without caching intermediate activations, allowing for
+ reduced memory at the expense of extra compute in the backward pass.
+ :param func: the function to evaluate.
+ :param inputs: the argument sequence to pass to `func`.
+ :param params: a sequence of parameters `func` depends on but does not
+ explicitly take as arguments.
+ :param flag: if False, disable gradient checkpointing.
+ """
+ if flag:
+ args = tuple(inputs) + tuple(params)
+ return CheckpointFunction.apply(func, len(inputs), *args)
+ else:
+ return func(*inputs)
+
+
+class CheckpointFunction(torch.autograd.Function):
+ @staticmethod
+ def forward(ctx, run_function, length, *args):
+ ctx.run_function = run_function
+ ctx.input_tensors = list(args[:length])
+ ctx.input_params = list(args[length:])
+
+ with torch.no_grad():
+ output_tensors = ctx.run_function(*ctx.input_tensors)
+ return output_tensors
+
+ @staticmethod
+ def backward(ctx, *output_grads):
+ ctx.input_tensors = [x.detach().requires_grad_(True) for x in ctx.input_tensors]
+ with torch.enable_grad():
+ # Fixes a bug where the first op in run_function modifies the
+ # Tensor storage in place, which is not allowed for detach()'d
+ # Tensors.
+ shallow_copies = [x.view_as(x) for x in ctx.input_tensors]
+ output_tensors = ctx.run_function(*shallow_copies)
+ input_grads = torch.autograd.grad(
+ output_tensors,
+ ctx.input_tensors + ctx.input_params,
+ output_grads,
+ allow_unused=True,
+ )
+ del ctx.input_tensors
+ del ctx.input_params
+ del output_tensors
+ return (None, None) + input_grads
+
+
+def timestep_embedding(timesteps, dim, max_period=10000, repeat_only=False):
+ """
+ Create sinusoidal timestep embeddings.
+ :param timesteps: a 1-D Tensor of N indices, one per batch element.
+ These may be fractional.
+ :param dim: the dimension of the output.
+ :param max_period: controls the minimum frequency of the embeddings.
+ :return: an [N x dim] Tensor of positional embeddings.
+ """
+ if not repeat_only:
+ half = dim // 2
+ freqs = torch.exp(-math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half).to(device=timesteps.device)
+ args = timesteps[:, None].float() * freqs[None]
+ embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+ if dim % 2:
+ embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
+ else:
+ embedding = repeat(timesteps, 'b -> b d', d=dim)
+ return embedding
+
+
+def zero_module(module):
+ """
+ Zero out the parameters of a module and return it.
+ """
+ for p in module.parameters():
+ p.detach().zero_()
+ return module
+
+
+def scale_module(module, scale):
+ """
+ Scale the parameters of a module and return it.
+ """
+ for p in module.parameters():
+ p.detach().mul_(scale)
+ return module
+
+
+def mean_flat(tensor):
+ """
+ Take the mean over all non-batch dimensions.
+ """
+ return tensor.mean(dim=list(range(1, len(tensor.shape))))
+
+
+def normalization(channels):
+ """
+ Make a standard normalization layer.
+ :param channels: number of input channels.
+ :return: an nn.Module for normalization.
+ """
+ return GroupNorm32(32, channels)
+
+
+# PyTorch 1.7 has SiLU, but we support PyTorch 1.5.
+class SiLU(nn.Module):
+ def forward(self, x):
+ return x * torch.sigmoid(x)
+
+
+class GroupNorm32(nn.GroupNorm):
+ def forward(self, x):
+ return super().forward(x.float()).type(x.dtype)
+
+
+def conv_nd(dims, *args, cconv=False, **kwargs):
+ """
+ Create a 1D, 2D, or 3D convolution module.
+ """
+ if dims == 1:
+ return nn.Conv1d(*args, **kwargs)
+ elif dims == 2:
+ if cconv:
+ return CircularConv2d(*args, **kwargs)
+ else:
+ return nn.Conv2d(*args, **kwargs)
+ elif dims == 3:
+ return nn.Conv3d(*args, **kwargs)
+ raise ValueError(f"unsupported dimensions: {dims}")
+
+
+def linear(*args, **kwargs):
+ """
+ Create a linear module.
+ """
+ return nn.Linear(*args, **kwargs)
+
+
+def avg_pool_nd(dims, *args, **kwargs):
+ """
+ Create a 1D, 2D, or 3D average pooling module.
+ """
+ if dims == 1:
+ return nn.AvgPool1d(*args, **kwargs)
+ elif dims == 2:
+ return nn.AvgPool2d(*args, **kwargs)
+ elif dims == 3:
+ return nn.AvgPool3d(*args, **kwargs)
+ raise ValueError(f"unsupported dimensions: {dims}")
+
+
+class HybridConditioner(nn.Module):
+
+ def __init__(self, c_concat_config, c_crossattn_config):
+ super().__init__()
+ self.concat_conditioner = instantiate_from_config(c_concat_config)
+ self.crossattn_conditioner = instantiate_from_config(c_crossattn_config)
+
+ def forward(self, c_concat, c_crossattn):
+ c_concat = self.concat_conditioner(c_concat)
+ c_crossattn = self.crossattn_conditioner(c_crossattn)
+ return {'c_concat': [c_concat], 'c_crossattn': [c_crossattn]}
+
+
+def noise_like(shape, device, repeat=False):
+ repeat_noise = lambda: torch.randn((1, *shape[1:]), device=device).repeat(shape[0], *((1,) * (len(shape) - 1)))
+ noise = lambda: torch.randn(shape, device=device)
+ return repeat_noise() if repeat else noise()
\ No newline at end of file
diff --git a/lidm/modules/diffusion/__init__.py b/lidm/modules/diffusion/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/lidm/modules/diffusion/model_ldm.py b/lidm/modules/diffusion/model_ldm.py
new file mode 100644
index 0000000000000000000000000000000000000000..5374a1142212bf9292a0fc3fbbfe7ff772aa8d0e
--- /dev/null
+++ b/lidm/modules/diffusion/model_ldm.py
@@ -0,0 +1,817 @@
+# pytorch_diffusion + derived encoder decoder
+import math
+import torch
+import torch.nn as nn
+import numpy as np
+from einops import rearrange
+
+from ...utils.misc_utils import instantiate_from_config
+from ...modules.attention import LinearAttention
+
+
+def get_timestep_embedding(timesteps, embedding_dim):
+ """
+ This matches the implementation in Denoising Diffusion Probabilistic Models:
+ From Fairseq.
+ Build sinusoidal embeddings.
+ This matches the implementation in tensor2tensor, but differs slightly
+ from the description in Section 3.5 of "Attention Is All You Need".
+ """
+ assert len(timesteps.shape) == 1
+
+ half_dim = embedding_dim // 2
+ emb = math.log(10000) / (half_dim - 1)
+ emb = torch.exp(torch.arange(half_dim, dtype=torch.float32) * -emb)
+ emb = emb.to(device=timesteps.device)
+ emb = timesteps.float()[:, None] * emb[None, :]
+ emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1)
+ if embedding_dim % 2 == 1: # zero pad
+ emb = torch.nn.functional.pad(emb, (0, 1, 0, 0))
+ return emb
+
+
+def nonlinearity(x):
+ # swish
+ return x * torch.sigmoid(x)
+
+
+def Normalize(in_channels, num_groups=32):
+ return torch.nn.GroupNorm(num_groups=num_groups, num_channels=in_channels, eps=1e-6, affine=True)
+
+
+class Upsample(nn.Module):
+ def __init__(self, in_channels, with_conv):
+ super().__init__()
+ self.with_conv = with_conv
+ if self.with_conv:
+ self.conv = torch.nn.Conv2d(in_channels,
+ in_channels,
+ kernel_size=3,
+ stride=1,
+ padding=1)
+
+ def forward(self, x):
+ x = torch.nn.functional.interpolate(x, scale_factor=2.0, mode="nearest")
+ if self.with_conv:
+ x = self.conv(x)
+ return x
+
+
+class Downsample(nn.Module):
+ def __init__(self, in_channels, with_conv):
+ super().__init__()
+ self.with_conv = with_conv
+ if self.with_conv:
+ # no asymmetric padding in torch conv, must do it ourselves
+ self.conv = torch.nn.Conv2d(in_channels,
+ in_channels,
+ kernel_size=3,
+ stride=2,
+ padding=0)
+
+ def forward(self, x):
+ if self.with_conv:
+ pad = (0, 1, 0, 1)
+ x = torch.nn.functional.pad(x, pad, mode="constant", value=0)
+ x = self.conv(x)
+ else:
+ x = torch.nn.functional.avg_pool2d(x, kernel_size=2, stride=2)
+ return x
+
+
+class ResnetBlock(nn.Module):
+ def __init__(self, *, in_channels, out_channels=None, conv_shortcut=False,
+ dropout, temb_channels=512):
+ super().__init__()
+ self.in_channels = in_channels
+ out_channels = in_channels if out_channels is None else out_channels
+ self.out_channels = out_channels
+ self.use_conv_shortcut = conv_shortcut
+
+ self.norm1 = Normalize(in_channels)
+ self.conv1 = torch.nn.Conv2d(in_channels,
+ out_channels,
+ kernel_size=3,
+ stride=1,
+ padding=1)
+ if temb_channels > 0:
+ self.temb_proj = torch.nn.Linear(temb_channels,
+ out_channels)
+ self.norm2 = Normalize(out_channels)
+ self.dropout = torch.nn.Dropout(dropout)
+ self.conv2 = torch.nn.Conv2d(out_channels,
+ out_channels,
+ kernel_size=3,
+ stride=1,
+ padding=1)
+ if self.in_channels != self.out_channels:
+ if self.use_conv_shortcut:
+ self.conv_shortcut = torch.nn.Conv2d(in_channels,
+ out_channels,
+ kernel_size=3,
+ stride=1,
+ padding=1)
+ else:
+ self.nin_shortcut = torch.nn.Conv2d(in_channels,
+ out_channels,
+ kernel_size=1,
+ stride=1,
+ padding=0)
+
+ def forward(self, x, temb):
+ h = x
+ h = self.norm1(h)
+ h = nonlinearity(h)
+ h = self.conv1(h)
+
+ if temb is not None:
+ h = h + self.temb_proj(nonlinearity(temb))[:, :, None, None]
+
+ h = self.norm2(h)
+ h = nonlinearity(h)
+ h = self.dropout(h)
+ h = self.conv2(h)
+
+ if self.in_channels != self.out_channels:
+ if self.use_conv_shortcut:
+ x = self.conv_shortcut(x)
+ else:
+ x = self.nin_shortcut(x)
+
+ return x + h
+
+
+class LinAttnBlock(LinearAttention):
+ """to match AttnBlock usage"""
+
+ def __init__(self, in_channels):
+ super().__init__(dim=in_channels, heads=1, dim_head=in_channels)
+
+
+class AttnBlock(nn.Module):
+ def __init__(self, in_channels):
+ super().__init__()
+ self.in_channels = in_channels
+
+ self.norm = Normalize(in_channels)
+ self.q = torch.nn.Conv2d(in_channels,
+ in_channels,
+ kernel_size=1,
+ stride=1,
+ padding=0)
+ self.k = torch.nn.Conv2d(in_channels,
+ in_channels,
+ kernel_size=1,
+ stride=1,
+ padding=0)
+ self.v = torch.nn.Conv2d(in_channels,
+ in_channels,
+ kernel_size=1,
+ stride=1,
+ padding=0)
+ self.proj_out = torch.nn.Conv2d(in_channels,
+ in_channels,
+ kernel_size=1,
+ stride=1,
+ padding=0)
+
+ def forward(self, x):
+ h_ = x
+ h_ = self.norm(h_)
+ q = self.q(h_)
+ k = self.k(h_)
+ v = self.v(h_)
+
+ # compute attention
+ b, c, h, w = q.shape
+ q = q.reshape(b, c, h * w)
+ q = q.permute(0, 2, 1) # b,hw,c
+ k = k.reshape(b, c, h * w) # b,c,hw
+ w_ = torch.bmm(q, k) # b,hw,hw w[b,i,j]=sum_c q[b,i,c]k[b,c,j]
+ w_ = w_ * (int(c) ** (-0.5))
+ w_ = torch.nn.functional.softmax(w_, dim=2)
+
+ # attend to values
+ v = v.reshape(b, c, h * w)
+ w_ = w_.permute(0, 2, 1) # b,hw,hw (first hw of k, second of q)
+ h_ = torch.bmm(v, w_) # b, c,hw (hw of q) h_[b,c,j] = sum_i v[b,c,i] w_[b,i,j]
+ h_ = h_.reshape(b, c, h, w)
+
+ h_ = self.proj_out(h_)
+
+ return x + h_
+
+
+def make_attn(in_channels, attn_type="vanilla"):
+ assert attn_type in ["vanilla", "linear", "none"], f'attn_type {attn_type} unknown'
+ # print(f"making attention of type '{attn_type}' with {in_channels} in_channels")
+ if attn_type == "vanilla":
+ return AttnBlock(in_channels)
+ elif attn_type == "none":
+ return nn.Identity(in_channels)
+ else:
+ return LinAttnBlock(in_channels)
+
+
+class Model(nn.Module):
+ def __init__(self, *, ch, out_ch, ch_mult=(1, 2, 4, 8), num_res_blocks,
+ attn_levels, dropout=0.0, resamp_with_conv=True, in_channels,
+ use_timestep=True, use_linear_attn=False, attn_type="vanilla"):
+ super().__init__()
+ if use_linear_attn: attn_type = "linear"
+ self.ch = ch
+ self.temb_ch = self.ch * 4
+ self.num_resolutions = len(ch_mult)
+ self.num_res_blocks = num_res_blocks
+ self.in_channels = in_channels
+
+ self.use_timestep = use_timestep
+ if self.use_timestep:
+ # timestep embedding
+ self.temb = nn.Module()
+ self.temb.dense = nn.ModuleList([
+ torch.nn.Linear(self.ch,
+ self.temb_ch),
+ torch.nn.Linear(self.temb_ch,
+ self.temb_ch),
+ ])
+
+ # downsampling
+ self.conv_in = torch.nn.Conv2d(in_channels,
+ self.ch,
+ kernel_size=3,
+ stride=1,
+ padding=1)
+
+ in_ch_mult = (1,) + tuple(ch_mult)
+ self.down = nn.ModuleList()
+ for i_level in range(self.num_resolutions):
+ block = nn.ModuleList()
+ attn = nn.ModuleList()
+ block_in = ch * in_ch_mult[i_level]
+ block_out = ch * ch_mult[i_level]
+ for i_block in range(self.num_res_blocks):
+ block.append(ResnetBlock(in_channels=block_in,
+ out_channels=block_out,
+ temb_channels=self.temb_ch,
+ dropout=dropout))
+ block_in = block_out
+ if i_level in attn_levels:
+ attn.append(make_attn(block_in, attn_type=attn_type))
+ down = nn.Module()
+ down.block = block
+ down.attn = attn
+ if i_level != self.num_resolutions - 1:
+ down.downsample = Downsample(block_in, resamp_with_conv)
+ self.down.append(down)
+
+ # middle
+ self.mid = nn.Module()
+ self.mid.block_1 = ResnetBlock(in_channels=block_in,
+ out_channels=block_in,
+ temb_channels=self.temb_ch,
+ dropout=dropout)
+ self.mid.attn_1 = make_attn(block_in, attn_type=attn_type)
+ self.mid.block_2 = ResnetBlock(in_channels=block_in,
+ out_channels=block_in,
+ temb_channels=self.temb_ch,
+ dropout=dropout)
+
+ # upsampling
+ self.up = nn.ModuleList()
+ for i_level in reversed(range(self.num_resolutions)):
+ block = nn.ModuleList()
+ attn = nn.ModuleList()
+ block_out = ch * ch_mult[i_level]
+ skip_in = ch * ch_mult[i_level]
+ for i_block in range(self.num_res_blocks + 1):
+ if i_block == self.num_res_blocks:
+ skip_in = ch * in_ch_mult[i_level]
+ block.append(ResnetBlock(in_channels=block_in + skip_in,
+ out_channels=block_out,
+ temb_channels=self.temb_ch,
+ dropout=dropout))
+ block_in = block_out
+ if i_level in attn_levels:
+ attn.append(make_attn(block_in, attn_type=attn_type))
+ up = nn.Module()
+ up.block = block
+ up.attn = attn
+ if i_level != 0:
+ up.upsample = Upsample(block_in, resamp_with_conv)
+ self.up.insert(0, up) # prepend to get consistent order
+
+ # end
+ self.norm_out = Normalize(block_in)
+ self.conv_out = torch.nn.Conv2d(block_in,
+ out_ch,
+ kernel_size=3,
+ stride=1,
+ padding=1)
+
+ def forward(self, x, t=None, context=None):
+ # assert x.shape[2] == x.shape[3] == self.resolution
+ if context is not None:
+ # assume aligned context, cat along channel axis
+ x = torch.cat((x, context), dim=1)
+ if self.use_timestep:
+ # timestep embedding
+ assert t is not None
+ temb = get_timestep_embedding(t, self.ch)
+ temb = self.temb.dense[0](temb)
+ temb = nonlinearity(temb)
+ temb = self.temb.dense[1](temb)
+ else:
+ temb = None
+
+ # downsampling
+ hs = [self.conv_in(x)]
+ for i_level in range(self.num_resolutions):
+ for i_block in range(self.num_res_blocks):
+ h = self.down[i_level].block[i_block](hs[-1], temb)
+ if len(self.down[i_level].attn) > 0:
+ h = self.down[i_level].attn[i_block](h)
+ hs.append(h)
+ if i_level != self.num_resolutions - 1:
+ hs.append(self.down[i_level].downsample(hs[-1]))
+
+ # middle
+ h = hs[-1]
+ h = self.mid.block_1(h, temb)
+ h = self.mid.attn_1(h)
+ h = self.mid.block_2(h, temb)
+
+ # upsampling
+ for i_level in reversed(range(self.num_resolutions)):
+ for i_block in range(self.num_res_blocks + 1):
+ h = self.up[i_level].block[i_block](
+ torch.cat([h, hs.pop()], dim=1), temb)
+ if len(self.up[i_level].attn) > 0:
+ h = self.up[i_level].attn[i_block](h)
+ if i_level != 0:
+ h = self.up[i_level].upsample(h)
+
+ # end
+ h = self.norm_out(h)
+ h = nonlinearity(h)
+ h = self.conv_out(h)
+ return h
+
+ def get_last_layer(self):
+ return self.conv_out.weight
+
+
+class Encoder(nn.Module):
+ def __init__(self, *, ch, out_ch, ch_mult=(1, 2, 4, 8), num_res_blocks,
+ attn_levels, dropout=0.0, resamp_with_conv=True, in_channels,
+ z_channels, double_z=True, use_linear_attn=False, attn_type="vanilla",
+ **ignore_kwargs):
+ super().__init__()
+ if use_linear_attn: attn_type = "linear"
+ self.ch = ch
+ self.temb_ch = 0
+ self.num_resolutions = len(ch_mult)
+ self.num_res_blocks = num_res_blocks
+ self.in_channels = in_channels
+
+ # downsampling
+ self.conv_in = torch.nn.Conv2d(in_channels,
+ self.ch,
+ kernel_size=3,
+ stride=1,
+ padding=1)
+
+ in_ch_mult = (1,) + tuple(ch_mult)
+ self.in_ch_mult = in_ch_mult
+ self.down = nn.ModuleList()
+ for i_level in range(self.num_resolutions):
+ block = nn.ModuleList()
+ attn = nn.ModuleList()
+ block_in = ch * in_ch_mult[i_level]
+ block_out = ch * ch_mult[i_level]
+ for i_block in range(self.num_res_blocks):
+ block.append(ResnetBlock(in_channels=block_in,
+ out_channels=block_out,
+ temb_channels=self.temb_ch,
+ dropout=dropout))
+ block_in = block_out
+ if i_level in attn_levels:
+ attn.append(make_attn(block_in, attn_type=attn_type))
+ down = nn.Module()
+ down.block = block
+ down.attn = attn
+ if i_level != self.num_resolutions - 1:
+ down.downsample = Downsample(block_in, resamp_with_conv)
+ self.down.append(down)
+
+ # middle
+ self.mid = nn.Module()
+ self.mid.block_1 = ResnetBlock(in_channels=block_in,
+ out_channels=block_in,
+ temb_channels=self.temb_ch,
+ dropout=dropout)
+ self.mid.attn_1 = make_attn(block_in, attn_type=attn_type)
+ self.mid.block_2 = ResnetBlock(in_channels=block_in,
+ out_channels=block_in,
+ temb_channels=self.temb_ch,
+ dropout=dropout)
+
+ # end
+ self.norm_out = Normalize(block_in)
+ self.conv_out = torch.nn.Conv2d(block_in,
+ 2 * z_channels if double_z else z_channels,
+ kernel_size=3,
+ stride=1,
+ padding=1)
+
+ def forward(self, x):
+ # timestep embedding
+ temb = None
+ # downsampling
+ hs = [self.conv_in(x)]
+ for i_level in range(self.num_resolutions):
+ for i_block in range(self.num_res_blocks):
+ h = self.down[i_level].block[i_block](hs[-1], temb)
+ if len(self.down[i_level].attn) > 0:
+ h = self.down[i_level].attn[i_block](h)
+ hs.append(h)
+ if i_level != self.num_resolutions - 1:
+ hs.append(self.down[i_level].downsample(hs[-1]))
+
+ # middle
+ h = hs[-1]
+ h = self.mid.block_1(h, temb)
+ h = self.mid.attn_1(h)
+ h = self.mid.block_2(h, temb)
+
+ # end
+ h = self.norm_out(h)
+ h = nonlinearity(h)
+ h = self.conv_out(h)
+ return h
+
+
+class Decoder(nn.Module):
+ def __init__(self, *, ch, out_ch, ch_mult=(1, 2, 4, 8), num_res_blocks,
+ attn_levels, dropout=0.0, resamp_with_conv=True, in_channels,
+ z_channels, give_pre_end=False, tanh_out=False, use_linear_attn=False,
+ attn_type="vanilla", **ignorekwargs):
+ super().__init__()
+ if use_linear_attn: attn_type = "linear"
+ self.ch = ch
+ self.temb_ch = 0
+ self.num_resolutions = len(ch_mult)
+ self.num_res_blocks = num_res_blocks
+ self.in_channels = in_channels
+ self.give_pre_end = give_pre_end
+ self.tanh_out = tanh_out
+
+ # compute in_ch_mult, block_in and curr_res at lowest res
+ in_ch_mult = (1,) + tuple(ch_mult)
+ block_in = ch * ch_mult[self.num_resolutions - 1]
+
+ # z to block_in
+ self.conv_in = torch.nn.Conv2d(z_channels,
+ block_in,
+ kernel_size=3,
+ stride=1,
+ padding=1)
+
+ # middle
+ self.mid = nn.Module()
+ self.mid.block_1 = ResnetBlock(in_channels=block_in,
+ out_channels=block_in,
+ temb_channels=self.temb_ch,
+ dropout=dropout)
+ self.mid.attn_1 = make_attn(block_in, attn_type=attn_type)
+ self.mid.block_2 = ResnetBlock(in_channels=block_in,
+ out_channels=block_in,
+ temb_channels=self.temb_ch,
+ dropout=dropout)
+
+ # upsampling
+ self.up = nn.ModuleList()
+ for i_level in reversed(range(self.num_resolutions)):
+ block = nn.ModuleList()
+ attn = nn.ModuleList()
+ block_out = ch * ch_mult[i_level]
+ for i_block in range(self.num_res_blocks + 1):
+ block.append(ResnetBlock(in_channels=block_in,
+ out_channels=block_out,
+ temb_channels=self.temb_ch,
+ dropout=dropout))
+ block_in = block_out
+ if i_level in attn_levels:
+ attn.append(make_attn(block_in, attn_type=attn_type))
+ up = nn.Module()
+ up.block = block
+ up.attn = attn
+ if i_level != 0:
+ up.upsample = Upsample(block_in, resamp_with_conv)
+ self.up.insert(0, up) # prepend to get consistent order
+
+ # end
+ self.norm_out = Normalize(block_in)
+ self.conv_out = torch.nn.Conv2d(block_in,
+ out_ch,
+ kernel_size=3,
+ stride=1,
+ padding=1)
+
+ def forward(self, z):
+ self.last_z_shape = z.shape
+
+ # timestep embedding
+ temb = None
+
+ # z to block_in
+ h = self.conv_in(z)
+
+ # middle
+ h = self.mid.block_1(h, temb)
+ h = self.mid.attn_1(h)
+ h = self.mid.block_2(h, temb)
+
+ # upsampling
+ for i_level in reversed(range(self.num_resolutions)):
+ for i_block in range(self.num_res_blocks + 1):
+ h = self.up[i_level].block[i_block](h, temb)
+ if len(self.up[i_level].attn) > 0:
+ h = self.up[i_level].attn[i_block](h)
+ if i_level != 0:
+ h = self.up[i_level].upsample(h)
+
+ # end
+ if self.give_pre_end:
+ return h
+
+ h = self.norm_out(h)
+ h = nonlinearity(h)
+ h = self.conv_out(h)
+ if self.tanh_out:
+ h = torch.tanh(h)
+ return h
+
+
+class SimpleDecoder(nn.Module):
+ def __init__(self, in_channels, out_channels, *args, **kwargs):
+ super().__init__()
+ self.model = nn.ModuleList([nn.Conv2d(in_channels, in_channels, 1),
+ ResnetBlock(in_channels=in_channels,
+ out_channels=2 * in_channels,
+ temb_channels=0, dropout=0.0),
+ ResnetBlock(in_channels=2 * in_channels,
+ out_channels=4 * in_channels,
+ temb_channels=0, dropout=0.0),
+ ResnetBlock(in_channels=4 * in_channels,
+ out_channels=2 * in_channels,
+ temb_channels=0, dropout=0.0),
+ nn.Conv2d(2 * in_channels, in_channels, 1),
+ Upsample(in_channels, with_conv=True)])
+ # end
+ self.norm_out = Normalize(in_channels)
+ self.conv_out = torch.nn.Conv2d(in_channels,
+ out_channels,
+ kernel_size=3,
+ stride=1,
+ padding=1)
+
+ def forward(self, x):
+ for i, layer in enumerate(self.model):
+ if i in [1, 2, 3]:
+ x = layer(x, None)
+ else:
+ x = layer(x)
+
+ h = self.norm_out(x)
+ h = nonlinearity(h)
+ x = self.conv_out(h)
+ return x
+
+
+class UpsampleDecoder(nn.Module):
+ def __init__(self, in_channels, out_channels, ch, num_res_blocks,
+ ch_mult=(2, 2), dropout=0.0):
+ super().__init__()
+ # upsampling
+ self.temb_ch = 0
+ self.num_resolutions = len(ch_mult)
+ block_in = in_channels
+ self.res_blocks = nn.ModuleList()
+ self.upsample_blocks = nn.ModuleList()
+ for i_level in range(self.num_resolutions):
+ res_block = []
+ block_out = ch * ch_mult[i_level]
+ for i_block in range(self.num_res_blocks + 1):
+ res_block.append(ResnetBlock(in_channels=block_in,
+ out_channels=block_out,
+ temb_channels=self.temb_ch,
+ dropout=dropout))
+ block_in = block_out
+ self.res_blocks.append(nn.ModuleList(res_block))
+ if i_level != self.num_resolutions - 1:
+ self.upsample_blocks.append(Upsample(block_in, True))
+
+ # end
+ self.norm_out = Normalize(block_in)
+ self.conv_out = torch.nn.Conv2d(block_in,
+ out_channels,
+ kernel_size=3,
+ stride=1,
+ padding=1)
+
+ def forward(self, x):
+ # upsampling
+ h = x
+ for k, i_level in enumerate(range(self.num_resolutions)):
+ for i_block in range(self.num_res_blocks + 1):
+ h = self.res_blocks[i_level][i_block](h, None)
+ if i_level != self.num_resolutions - 1:
+ h = self.upsample_blocks[k](h)
+ h = self.norm_out(h)
+ h = nonlinearity(h)
+ h = self.conv_out(h)
+ return h
+
+
+class LatentRescaler(nn.Module):
+ def __init__(self, factor, in_channels, mid_channels, out_channels, depth=2):
+ super().__init__()
+ # residual block, interpolate, residual block
+ self.factor = factor
+ self.conv_in = nn.Conv2d(in_channels,
+ mid_channels,
+ kernel_size=3,
+ stride=1,
+ padding=1)
+ self.res_block1 = nn.ModuleList([ResnetBlock(in_channels=mid_channels,
+ out_channels=mid_channels,
+ temb_channels=0,
+ dropout=0.0) for _ in range(depth)])
+ self.attn = AttnBlock(mid_channels)
+ self.res_block2 = nn.ModuleList([ResnetBlock(in_channels=mid_channels,
+ out_channels=mid_channels,
+ temb_channels=0,
+ dropout=0.0) for _ in range(depth)])
+
+ self.conv_out = nn.Conv2d(mid_channels,
+ out_channels,
+ kernel_size=1,
+ )
+
+ def forward(self, x):
+ x = self.conv_in(x)
+ for block in self.res_block1:
+ x = block(x, None)
+ x = torch.nn.functional.interpolate(x, size=(
+ int(round(x.shape[2] * self.factor)), int(round(x.shape[3] * self.factor))))
+ x = self.attn(x)
+ for block in self.res_block2:
+ x = block(x, None)
+ x = self.conv_out(x)
+ return x
+
+
+class MergedRescaleEncoder(nn.Module):
+ def __init__(self, in_channels, ch, out_ch, num_res_blocks,
+ attn_levels, dropout=0.0, resamp_with_conv=True,
+ ch_mult=(1, 2, 4, 8), rescale_factor=1.0, rescale_module_depth=1):
+ super().__init__()
+ intermediate_chn = ch * ch_mult[-1]
+ self.encoder = Encoder(in_channels=in_channels, num_res_blocks=num_res_blocks, ch=ch, ch_mult=ch_mult,
+ z_channels=intermediate_chn, double_z=False,
+ attn_levels=attn_levels, dropout=dropout, resamp_with_conv=resamp_with_conv,
+ out_ch=None)
+ self.rescaler = LatentRescaler(factor=rescale_factor, in_channels=intermediate_chn,
+ mid_channels=intermediate_chn, out_channels=out_ch, depth=rescale_module_depth)
+
+ def forward(self, x):
+ x = self.encoder(x)
+ x = self.rescaler(x)
+ return x
+
+
+class MergedRescaleDecoder(nn.Module):
+ def __init__(self, z_channels, out_ch, num_res_blocks, attn_levels, ch, ch_mult=(1, 2, 4, 8),
+ dropout=0.0, resamp_with_conv=True, rescale_factor=1.0, rescale_module_depth=1):
+ super().__init__()
+ tmp_chn = z_channels * ch_mult[-1]
+ self.decoder = Decoder(out_ch=out_ch, z_channels=tmp_chn, attn_levels=attn_levels, dropout=dropout,
+ resamp_with_conv=resamp_with_conv, in_channels=None, num_res_blocks=num_res_blocks,
+ ch_mult=ch_mult, ch=ch)
+ self.rescaler = LatentRescaler(factor=rescale_factor, in_channels=z_channels, mid_channels=tmp_chn,
+ out_channels=tmp_chn, depth=rescale_module_depth)
+
+ def forward(self, x):
+ x = self.rescaler(x)
+ x = self.decoder(x)
+ return x
+
+
+class Upsampler(nn.Module):
+ def __init__(self, in_size, out_size, in_channels, out_channels, ch_mult=2):
+ super().__init__()
+ assert out_size >= in_size
+ num_blocks = int(np.log2(out_size // in_size)) + 1
+ factor_up = 1. + (out_size % in_size)
+ print(
+ f"Building {self.__class__.__name__} with in_size: {in_size} --> out_size {out_size} and factor {factor_up}")
+ self.rescaler = LatentRescaler(factor=factor_up, in_channels=in_channels, mid_channels=2 * in_channels,
+ out_channels=in_channels)
+ self.decoder = Decoder(out_ch=out_channels, z_channels=in_channels, num_res_blocks=2,
+ attn_levels=[], in_channels=None, ch=in_channels,
+ ch_mult=[ch_mult for _ in range(num_blocks)])
+
+ def forward(self, x):
+ x = self.rescaler(x)
+ x = self.decoder(x)
+ return x
+
+
+class Resize(nn.Module):
+ def __init__(self, in_channels=None, learned=False, mode="bilinear"):
+ super().__init__()
+ self.with_conv = learned
+ self.mode = mode
+ if self.with_conv:
+ print(f"Note: {self.__class__.__name} uses learned downsampling and will ignore the fixed {mode} mode")
+ raise NotImplementedError()
+ assert in_channels is not None
+ # no asymmetric padding in torch conv, must do it ourselves
+ self.conv = torch.nn.Conv2d(in_channels,
+ in_channels,
+ kernel_size=4,
+ stride=2,
+ padding=1)
+
+ def forward(self, x, scale_factor=1.0):
+ if scale_factor == 1.0:
+ return x
+ else:
+ x = torch.nn.functional.interpolate(x, mode=self.mode, align_corners=False, scale_factor=scale_factor)
+ return x
+
+
+class FirstStagePostProcessor(nn.Module):
+
+ def __init__(self, ch_mult: list, in_channels,
+ pretrained_model: nn.Module = None,
+ reshape=False,
+ n_channels=None,
+ dropout=0.,
+ pretrained_config=None):
+ super().__init__()
+ if pretrained_config is None:
+ assert pretrained_model is not None, 'Either "pretrained_model" or "pretrained_config" must not be None'
+ self.pretrained_model = pretrained_model
+ else:
+ assert pretrained_config is not None, 'Either "pretrained_model" or "pretrained_config" must not be None'
+ self.instantiate_pretrained(pretrained_config)
+
+ self.do_reshape = reshape
+
+ if n_channels is None:
+ n_channels = self.pretrained_model.encoder.ch
+
+ self.proj_norm = Normalize(in_channels, num_groups=in_channels // 2)
+ self.proj = nn.Conv2d(in_channels, n_channels, kernel_size=3,
+ stride=1, padding=1)
+
+ blocks = []
+ downs = []
+ ch_in = n_channels
+ for m in ch_mult:
+ blocks.append(ResnetBlock(in_channels=ch_in, out_channels=m * n_channels, dropout=dropout))
+ ch_in = m * n_channels
+ downs.append(Downsample(ch_in, with_conv=False))
+
+ self.model = nn.ModuleList(blocks)
+ self.downsampler = nn.ModuleList(downs)
+
+ def instantiate_pretrained(self, config):
+ model = instantiate_from_config(config)
+ self.pretrained_model = model.eval()
+ # self.pretrained_model.train = False
+ for param in self.pretrained_model.parameters():
+ param.requires_grad = False
+
+ @torch.no_grad()
+ def encode_with_pretrained(self, x):
+ c = self.pretrained_model.encode(x)
+ if isinstance(c, DiagonalGaussianDistribution):
+ c = c.mode()
+ return c
+
+ def forward(self, x):
+ z_fs = self.encode_with_pretrained(x)
+ z = self.proj_norm(z_fs)
+ z = self.proj(z)
+ z = nonlinearity(z)
+
+ for submodel, downmodel in zip(self.model, self.downsampler):
+ z = submodel(z, temb=None)
+ z = downmodel(z)
+
+ if self.do_reshape:
+ z = rearrange(z, 'b c h w -> b (h w) c')
+ return z
diff --git a/lidm/modules/diffusion/model_lidm.py b/lidm/modules/diffusion/model_lidm.py
new file mode 100644
index 0000000000000000000000000000000000000000..2803291c509d6d7c8bc5d5fb519fd7c46fd14707
--- /dev/null
+++ b/lidm/modules/diffusion/model_lidm.py
@@ -0,0 +1,681 @@
+# pytorch_diffusion + derived encoder decoder
+import math
+
+import torch
+import torch.nn as nn
+import numpy as np
+from einops import rearrange
+
+from ..basic import CircularConv2d
+from ...utils.misc_utils import instantiate_from_config
+from ...modules.attention import LinearAttention
+
+
+def get_timestep_embedding(timesteps, embedding_dim):
+ """
+ This matches the implementation in Denoising Diffusion Probabilistic Models:
+ From Fairseq.
+ Build sinusoidal embeddings.
+ This matches the implementation in tensor2tensor, but differs slightly
+ from the description in Section 3.5 of "Attention Is All You Need".
+ """
+ assert len(timesteps.shape) == 1
+
+ half_dim = embedding_dim // 2
+ emb = math.log(10000) / (half_dim - 1)
+ emb = torch.exp(torch.arange(half_dim, dtype=torch.float32) * -emb)
+ emb = emb.to(device=timesteps.device)
+ emb = timesteps.float()[:, None] * emb[None, :]
+ emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1)
+ if embedding_dim % 2 == 1: # zero pad
+ emb = torch.nn.functional.pad(emb, (0, 1, 0, 0))
+ return emb
+
+
+def nonlinearity(x):
+ # swish
+ return x * torch.sigmoid(x)
+
+
+def Normalize(in_channels, num_groups=32):
+ return torch.nn.GroupNorm(num_groups=num_groups, num_channels=in_channels, eps=1e-6, affine=True)
+
+
+UPSAMPLE_STRIDE2KERNEL_DICT = {(1, 2): (1, 5), (1, 4): (1, 7), (2, 1): (5, 1), (2, 2): (3, 3)}
+UPSAMPLE_STRIDE2PAD_DICT = {(1, 2): (2, 2, 0, 0), (1, 4): (3, 3, 0, 0), (2, 1): (0, 0, 2, 2), (2, 2): (1, 1, 1, 1)}
+
+
+class Upsample(nn.Module):
+ def __init__(self, in_channels, with_conv, stride):
+ super().__init__()
+ self.with_conv = with_conv
+ self.stride = stride
+ if self.with_conv:
+ k, p = UPSAMPLE_STRIDE2KERNEL_DICT[stride], UPSAMPLE_STRIDE2PAD_DICT[stride]
+ self.conv = CircularConv2d(in_channels, in_channels, kernel_size=k, padding=p)
+
+ def forward(self, x):
+ x = torch.nn.functional.interpolate(x, scale_factor=self.stride, mode='bilinear', align_corners=True)
+ if self.with_conv:
+ x = self.conv(x)
+ return x
+
+
+DOWNSAMPLE_STRIDE2KERNEL_DICT = {(1, 2): (3, 3), (1, 4): (3, 5), (2, 1): (3, 3), (2, 2): (3, 3)}
+DOWNSAMPLE_STRIDE2PAD_DICT = {(1, 2): (0, 1, 1, 1), (1, 4): (1, 1, 1, 1), (2, 1): (1, 1, 1, 1), (2, 2): (0, 1, 0, 1)}
+
+
+class Downsample(nn.Module):
+ def __init__(self, in_channels, with_conv, stride):
+ super().__init__()
+ self.with_conv = with_conv
+ self.stride = stride
+ if self.with_conv:
+ k, p = DOWNSAMPLE_STRIDE2KERNEL_DICT[stride], DOWNSAMPLE_STRIDE2PAD_DICT[stride]
+ self.conv = CircularConv2d(in_channels, in_channels, kernel_size=k, stride=stride, padding=p)
+
+ def forward(self, x):
+ if self.with_conv:
+ x = self.conv(x)
+ else:
+ x = torch.nn.functional.avg_pool2d(x, kernel_size=self.stride, stride=self.stride) # modified for lidar
+ return x
+
+
+UNIFORM_KERNEL2PAD_DICT = {(3, 3): (1, 1, 1, 1), (1, 4): (1, 2, 0, 0)}
+
+
+class ResnetBlock(nn.Module):
+ def __init__(self, *, in_channels, out_channels=None, kernel_size=(3, 3), conv_shortcut=False,
+ dropout, temb_channels=512):
+ super().__init__()
+ self.in_channels = in_channels
+ out_channels = in_channels if out_channels is None else out_channels
+ self.out_channels = out_channels
+ self.use_conv_shortcut = conv_shortcut
+ pad = UNIFORM_KERNEL2PAD_DICT[kernel_size]
+
+ self.norm1 = Normalize(in_channels)
+ self.conv1 = CircularConv2d(in_channels,
+ out_channels,
+ kernel_size=kernel_size,
+ stride=1,
+ padding=pad)
+ if temb_channels > 0:
+ self.temb_proj = torch.nn.Linear(temb_channels, out_channels)
+ self.norm2 = Normalize(out_channels)
+ self.dropout = torch.nn.Dropout(dropout)
+ self.conv2 = CircularConv2d(out_channels,
+ out_channels,
+ kernel_size=kernel_size,
+ stride=1,
+ padding=pad)
+ if self.in_channels != self.out_channels:
+ if self.use_conv_shortcut:
+ self.conv_shortcut = CircularConv2d(in_channels,
+ out_channels,
+ kernel_size=kernel_size,
+ stride=1,
+ padding=pad)
+ else:
+ self.nin_shortcut = torch.nn.Conv2d(in_channels,
+ out_channels,
+ kernel_size=1,
+ stride=1,
+ padding=0)
+
+ def forward(self, x, temb):
+ h = x
+ h = self.norm1(h)
+ h = nonlinearity(h)
+ h = self.conv1(h)
+
+ if temb is not None:
+ h = h + self.temb_proj(nonlinearity(temb))[:, :, None, None]
+
+ h = self.norm2(h)
+ h = nonlinearity(h)
+ h = self.dropout(h)
+ h = self.conv2(h)
+
+ if self.in_channels != self.out_channels:
+ if self.use_conv_shortcut:
+ x = self.conv_shortcut(x)
+ else:
+ x = self.nin_shortcut(x)
+
+ return x + h
+
+
+class LinAttnBlock(LinearAttention):
+ """to match AttnBlock usage"""
+
+ def __init__(self, in_channels):
+ super().__init__(dim=in_channels, heads=1, dim_head=in_channels)
+
+
+class AttnBlock(nn.Module):
+ def __init__(self, in_channels):
+ super().__init__()
+ self.in_channels = in_channels
+
+ self.norm = Normalize(in_channels)
+ self.q = torch.nn.Conv2d(in_channels,
+ in_channels,
+ kernel_size=1,
+ stride=1,
+ padding=0)
+ self.k = torch.nn.Conv2d(in_channels,
+ in_channels,
+ kernel_size=1,
+ stride=1,
+ padding=0)
+ self.v = torch.nn.Conv2d(in_channels,
+ in_channels,
+ kernel_size=1,
+ stride=1,
+ padding=0)
+ self.proj_out = torch.nn.Conv2d(in_channels,
+ in_channels,
+ kernel_size=1,
+ stride=1,
+ padding=0)
+
+ def forward(self, x):
+ h_ = x
+ h_ = self.norm(h_)
+ q = self.q(h_)
+ k = self.k(h_)
+ v = self.v(h_)
+
+ # compute attention
+ b, c, h, w = q.shape
+ q = q.reshape(b, c, h * w)
+ q = q.permute(0, 2, 1) # b,hw,c
+ k = k.reshape(b, c, h * w) # b,c,hw
+ w_ = torch.bmm(q, k) # b,hw,hw w[b,i,j]=sum_c q[b,i,c]k[b,c,j]
+ w_ = w_ * (int(c) ** (-0.5))
+ w_ = torch.nn.functional.softmax(w_, dim=2)
+
+ # attend to values
+ v = v.reshape(b, c, h * w)
+ w_ = w_.permute(0, 2, 1) # b,hw,hw (first hw of k, second of q)
+ h_ = torch.bmm(v, w_) # b, c,hw (hw of q) h_[b,c,j] = sum_i v[b,c,i] w_[b,i,j]
+ h_ = h_.reshape(b, c, h, w)
+
+ h_ = self.proj_out(h_)
+
+ return x + h_
+
+
+def make_attn(in_channels, attn_type="vanilla"):
+ assert attn_type in ["vanilla", "linear", "none"], f'attn_type {attn_type} unknown'
+ # print(f"making attention of type '{attn_type}' with {in_channels} in_channels")
+ if attn_type == "vanilla":
+ return AttnBlock(in_channels)
+ elif attn_type == "none":
+ return nn.Identity(in_channels)
+ else:
+ return LinAttnBlock(in_channels)
+
+
+class Encoder(nn.Module):
+ def __init__(self, *, ch, out_ch, ch_mult, strides, num_res_blocks,
+ attn_levels, dropout=0.0, resamp_with_conv=True, in_channels, z_channels,
+ double_z=True, use_linear_attn=False, attn_type="vanilla", use_mask=False,
+ **ignore_kwargs):
+ super().__init__()
+ if use_mask:
+ assert out_ch == in_channels + 1, 'Set "out_ch = out_ch + 1" for mask prediction.'
+ if use_linear_attn: attn_type = "linear"
+ self.ch = ch
+ self.temb_ch = 0
+ self.num_resolutions = len(ch_mult)
+ self.num_res_blocks = num_res_blocks
+ self.in_channels = in_channels
+
+ # downsampling
+ self.conv_in = CircularConv2d(in_channels,
+ self.ch,
+ kernel_size=3,
+ stride=1,
+ padding=1)
+ in_ch_mult = (1,) + tuple(ch_mult)
+ self.in_ch_mult = in_ch_mult
+ self.down = nn.ModuleList()
+ for i_level in range(self.num_resolutions):
+ block = nn.ModuleList()
+ attn = nn.ModuleList()
+ block_in = ch * in_ch_mult[i_level]
+ block_out = ch * ch_mult[i_level]
+ for i_block in range(self.num_res_blocks):
+ block.append(ResnetBlock(in_channels=block_in,
+ out_channels=block_out,
+ temb_channels=self.temb_ch,
+ dropout=dropout))
+ block_in = block_out
+ if i_level in attn_levels:
+ attn.append(make_attn(block_in, attn_type=attn_type))
+ down = nn.Module()
+ down.block = block
+ down.attn = attn
+ if i_level != self.num_resolutions - 1:
+ stride = tuple(strides[i_level])
+ down.downsample = Downsample(block_in, resamp_with_conv, stride)
+ self.down.append(down)
+
+ # middle
+ self.mid = nn.Module()
+ self.mid.block_1 = ResnetBlock(in_channels=block_in,
+ out_channels=block_in,
+ temb_channels=self.temb_ch,
+ dropout=dropout)
+ self.mid.attn_1 = make_attn(block_in, attn_type=attn_type)
+ self.mid.block_2 = ResnetBlock(in_channels=block_in,
+ out_channels=block_in,
+ temb_channels=self.temb_ch,
+ dropout=dropout)
+
+ # end
+ self.norm_out = Normalize(block_in)
+ self.conv_out = CircularConv2d(block_in,
+ 2 * z_channels if double_z else z_channels,
+ kernel_size=3,
+ stride=1,
+ padding=1)
+
+ def forward(self, x):
+ # timestep embedding
+ temb = None
+
+ # downsampling
+ hs = [self.conv_in(x)]
+ for i_level in range(self.num_resolutions):
+ for i_block in range(self.num_res_blocks):
+ h = self.down[i_level].block[i_block](hs[-1], temb)
+ if len(self.down[i_level].attn) > 0:
+ h = self.down[i_level].attn[i_block](h)
+ hs.append(h)
+ if i_level != self.num_resolutions - 1:
+ hs.append(self.down[i_level].downsample(hs[-1]))
+
+ # middle
+ h = hs[-1]
+ h = self.mid.block_1(h, temb)
+ h = self.mid.attn_1(h)
+ h = self.mid.block_2(h, temb)
+
+ # end
+ h = self.norm_out(h)
+ h = nonlinearity(h)
+ h = self.conv_out(h)
+ return h
+
+
+class Decoder(nn.Module):
+ def __init__(self, *, ch, out_ch, ch_mult, strides, num_res_blocks, attn_levels,
+ dropout=0.0, resamp_with_conv=True, in_channels, z_channels, give_pre_end=False,
+ tanh_out=False, use_linear_attn=False, attn_type="vanilla", use_mask=False,
+ **ignorekwargs):
+ super().__init__()
+ stride2kernel = {(2, 2): (3, 3), (1, 2): (1, 4)}
+ if use_linear_attn: attn_type = "linear"
+ self.ch = ch
+ self.temb_ch = 0
+ self.num_resolutions = len(ch_mult)
+ self.num_res_blocks = num_res_blocks
+ self.in_channels = in_channels
+ self.give_pre_end = give_pre_end
+ self.tanh_out = tanh_out
+
+ # compute in_ch_mult, block_in and curr_res at lowest res
+ block_in = ch * ch_mult[self.num_resolutions - 1]
+
+ # z to block_in
+ self.conv_in = CircularConv2d(z_channels,
+ block_in,
+ kernel_size=3,
+ stride=1,
+ padding=1)
+
+ # middle
+ self.mid = nn.Module()
+ self.mid.block_1 = ResnetBlock(in_channels=block_in,
+ out_channels=block_in,
+ temb_channels=self.temb_ch,
+ dropout=dropout)
+ self.mid.attn_1 = make_attn(block_in, attn_type=attn_type)
+ self.mid.block_2 = ResnetBlock(in_channels=block_in,
+ out_channels=block_in,
+ temb_channels=self.temb_ch,
+ dropout=dropout)
+
+ # upsampling
+ self.up = nn.ModuleList()
+ for i_level in reversed(range(self.num_resolutions)):
+ stride = tuple(strides[i_level - 1]) if i_level > 0 else None
+ kernel = stride2kernel[stride] if stride is not None else (1, 4)
+ block = nn.ModuleList()
+ attn = nn.ModuleList()
+ block_out = ch * ch_mult[i_level]
+ for i_block in range(self.num_res_blocks + 1):
+ block.append(ResnetBlock(in_channels=block_in,
+ out_channels=block_out,
+ kernel_size=kernel,
+ temb_channels=self.temb_ch,
+ dropout=dropout))
+ block_in = block_out
+ if i_level in attn_levels:
+ attn.append(make_attn(block_in, attn_type=attn_type))
+ up = nn.Module()
+ up.block = block
+ up.attn = attn
+ if stride is not None:
+ up.upsample = Upsample(block_in, resamp_with_conv, stride)
+ self.up.insert(0, up) # prepend to get consistent order
+
+ # end
+ self.norm_out = Normalize(block_in)
+ self.conv_out = CircularConv2d(block_in,
+ out_ch,
+ kernel_size=(1, 4),
+ stride=1,
+ padding=(1, 2, 0, 0))
+
+ def forward(self, z):
+ self.last_z_shape = z.shape
+
+ # timestep embedding
+ temb = None
+
+ # z to block_in
+ h = self.conv_in(z)
+
+ # middle
+ h = self.mid.block_1(h, temb)
+ h = self.mid.attn_1(h)
+ h = self.mid.block_2(h, temb)
+
+ # upsampling
+ for i_level in reversed(range(self.num_resolutions)):
+ for i_block in range(self.num_res_blocks + 1):
+ h = self.up[i_level].block[i_block](h, temb)
+ if len(self.up[i_level].attn) > 0:
+ h = self.up[i_level].attn[i_block](h)
+ if i_level != 0:
+ h = self.up[i_level].upsample(h)
+
+ # end
+ if self.give_pre_end:
+ return h
+
+ h = self.norm_out(h)
+ h = nonlinearity(h)
+ h = self.conv_out(h)
+ if self.tanh_out:
+ h = torch.tanh(h)
+ return h
+
+
+class SimpleDecoder(nn.Module):
+ def __init__(self, in_channels, out_channels, *args, **kwargs):
+ super().__init__()
+ self.model = nn.ModuleList([nn.Conv2d(in_channels, in_channels, 1),
+ ResnetBlock(in_channels=in_channels,
+ out_channels=2 * in_channels,
+ temb_channels=0, dropout=0.0),
+ ResnetBlock(in_channels=2 * in_channels,
+ out_channels=4 * in_channels,
+ temb_channels=0, dropout=0.0),
+ ResnetBlock(in_channels=4 * in_channels,
+ out_channels=2 * in_channels,
+ temb_channels=0, dropout=0.0),
+ nn.Conv2d(2 * in_channels, in_channels, 1),
+ Upsample(in_channels, with_conv=True)])
+ # end
+ self.norm_out = Normalize(in_channels)
+ self.conv_out = torch.nn.Conv2d(in_channels,
+ out_channels,
+ kernel_size=3,
+ stride=1,
+ padding=1)
+
+ def forward(self, x):
+ for i, layer in enumerate(self.model):
+ if i in [1, 2, 3]:
+ x = layer(x, None)
+ else:
+ x = layer(x)
+
+ h = self.norm_out(x)
+ h = nonlinearity(h)
+ x = self.conv_out(h)
+ return x
+
+
+class UpsampleDecoder(nn.Module):
+ def __init__(self, in_channels, out_channels, ch, num_res_blocks,
+ ch_mult=(2, 2), dropout=0.0):
+ super().__init__()
+ # upsampling
+ self.temb_ch = 0
+ self.num_resolutions = len(ch_mult)
+ self.num_res_blocks = num_res_blocks
+ block_in = in_channels
+ self.res_blocks = nn.ModuleList()
+ self.upsample_blocks = nn.ModuleList()
+ for i_level in range(self.num_resolutions):
+ res_block = []
+ block_out = ch * ch_mult[i_level]
+ for i_block in range(self.num_res_blocks + 1):
+ res_block.append(ResnetBlock(in_channels=block_in,
+ out_channels=block_out,
+ temb_channels=self.temb_ch,
+ dropout=dropout))
+ block_in = block_out
+ self.res_blocks.append(nn.ModuleList(res_block))
+ if i_level != self.num_resolutions - 1:
+ self.upsample_blocks.append(Upsample(block_in, True))
+
+ # end
+ self.norm_out = Normalize(block_in)
+ self.conv_out = torch.nn.Conv2d(block_in,
+ out_channels,
+ kernel_size=3,
+ stride=1,
+ padding=1)
+
+ def forward(self, x):
+ # upsampling
+ h = x
+ for k, i_level in enumerate(range(self.num_resolutions)):
+ for i_block in range(self.num_res_blocks + 1):
+ h = self.res_blocks[i_level][i_block](h, None)
+ if i_level != self.num_resolutions - 1:
+ h = self.upsample_blocks[k](h)
+ h = self.norm_out(h)
+ h = nonlinearity(h)
+ h = self.conv_out(h)
+ return h
+
+
+class LatentRescaler(nn.Module):
+ def __init__(self, factor, in_channels, mid_channels, out_channels, depth=2):
+ super().__init__()
+ # residual block, interpolate, residual block
+ self.factor = factor
+ self.conv_in = nn.Conv2d(in_channels,
+ mid_channels,
+ kernel_size=3,
+ stride=1,
+ padding=1)
+ self.res_block1 = nn.ModuleList([ResnetBlock(in_channels=mid_channels,
+ out_channels=mid_channels,
+ temb_channels=0,
+ dropout=0.0) for _ in range(depth)])
+ self.attn = AttnBlock(mid_channels)
+ self.res_block2 = nn.ModuleList([ResnetBlock(in_channels=mid_channels,
+ out_channels=mid_channels,
+ temb_channels=0,
+ dropout=0.0) for _ in range(depth)])
+
+ self.conv_out = nn.Conv2d(mid_channels,
+ out_channels,
+ kernel_size=1,
+ )
+
+ def forward(self, x):
+ x = self.conv_in(x)
+ for block in self.res_block1:
+ x = block(x, None)
+ x = torch.nn.functional.interpolate(x, size=(
+ int(round(x.shape[2] * self.factor)), int(round(x.shape[3] * self.factor))))
+ x = self.attn(x)
+ for block in self.res_block2:
+ x = block(x, None)
+ x = self.conv_out(x)
+ return x
+
+
+class MergedRescaleEncoder(nn.Module):
+ def __init__(self, in_channels, ch, out_ch, num_res_blocks,
+ attn_levels, dropout=0.0, resamp_with_conv=True,
+ ch_mult=(1, 2, 4, 8), rescale_factor=1.0, rescale_module_depth=1):
+ super().__init__()
+ intermediate_chn = ch * ch_mult[-1]
+ self.encoder = Encoder(in_channels=in_channels, num_res_blocks=num_res_blocks, ch=ch, ch_mult=ch_mult,
+ z_channels=intermediate_chn, double_z=False, attn_levels=attn_levels, dropout=dropout,
+ resamp_with_conv=resamp_with_conv, out_ch=None)
+ self.rescaler = LatentRescaler(factor=rescale_factor, in_channels=intermediate_chn,
+ mid_channels=intermediate_chn, out_channels=out_ch, depth=rescale_module_depth)
+
+ def forward(self, x):
+ x = self.encoder(x)
+ x = self.rescaler(x)
+ return x
+
+
+class MergedRescaleDecoder(nn.Module):
+ def __init__(self, z_channels, out_ch, num_res_blocks, attn_levels, ch, ch_mult=(1, 2, 4, 8),
+ dropout=0.0, resamp_with_conv=True, rescale_factor=1.0, rescale_module_depth=1):
+ super().__init__()
+ tmp_chn = z_channels * ch_mult[-1]
+ self.decoder = Decoder(out_ch=out_ch, z_channels=tmp_chn, attn_levels=attn_levels, dropout=dropout,
+ resamp_with_conv=resamp_with_conv, in_channels=None, num_res_blocks=num_res_blocks,
+ ch_mult=ch_mult, ch=ch)
+ self.rescaler = LatentRescaler(factor=rescale_factor, in_channels=z_channels, mid_channels=tmp_chn,
+ out_channels=tmp_chn, depth=rescale_module_depth)
+
+ def forward(self, x):
+ x = self.rescaler(x)
+ x = self.decoder(x)
+ return x
+
+
+class Upsampler(nn.Module):
+ def __init__(self, in_size, out_size, in_channels, out_channels, ch_mult=2):
+ super().__init__()
+ assert out_size >= in_size
+ num_blocks = int(np.log2(out_size // in_size)) + 1
+ factor_up = 1. + (out_size % in_size)
+ print(
+ f"Building {self.__class__.__name__} with in_size: {in_size} --> out_size {out_size} and factor {factor_up}")
+ self.rescaler = LatentRescaler(factor=factor_up, in_channels=in_channels, mid_channels=2 * in_channels,
+ out_channels=in_channels)
+ self.decoder = Decoder(out_ch=out_channels, z_channels=in_channels, num_res_blocks=2,
+ attn_levels=[], in_channels=None, ch=in_channels,
+ ch_mult=[ch_mult for _ in range(num_blocks)])
+
+ def forward(self, x):
+ x = self.rescaler(x)
+ x = self.decoder(x)
+ return x
+
+
+class Resize(nn.Module):
+ def __init__(self, in_channels=None, learned=False, mode="bilinear"):
+ super().__init__()
+ self.with_conv = learned
+ self.mode = mode
+ if self.with_conv:
+ print(f"Note: {self.__class__.__name} uses learned downsampling and will ignore the fixed {mode} mode")
+ raise NotImplementedError()
+ assert in_channels is not None
+ # no asymmetric padding in torch conv, must do it ourselves
+ self.conv = torch.nn.Conv2d(in_channels,
+ in_channels,
+ kernel_size=4,
+ stride=2,
+ padding=1)
+
+ def forward(self, x, scale_factor=1.0):
+ if scale_factor == 1.0:
+ return x
+ else:
+ x = torch.nn.functional.interpolate(x, mode=self.mode, align_corners=False, scale_factor=scale_factor)
+ return x
+
+
+class FirstStagePostProcessor(nn.Module):
+
+ def __init__(self, ch_mult: list, in_channels,
+ pretrained_model: nn.Module = None,
+ reshape=False,
+ n_channels=None,
+ dropout=0.,
+ pretrained_config=None):
+ super().__init__()
+ if pretrained_config is None:
+ assert pretrained_model is not None, 'Either "pretrained_model" or "pretrained_config" must not be None'
+ self.pretrained_model = pretrained_model
+ else:
+ assert pretrained_config is not None, 'Either "pretrained_model" or "pretrained_config" must not be None'
+ self.instantiate_pretrained(pretrained_config)
+
+ self.do_reshape = reshape
+
+ if n_channels is None:
+ n_channels = self.pretrained_model.encoder.ch
+
+ self.proj_norm = Normalize(in_channels, num_groups=in_channels // 2)
+ self.proj = nn.Conv2d(in_channels, n_channels, kernel_size=3,
+ stride=1, padding=1)
+
+ blocks = []
+ downs = []
+ ch_in = n_channels
+ for m in ch_mult:
+ blocks.append(ResnetBlock(in_channels=ch_in, out_channels=m * n_channels, dropout=dropout))
+ ch_in = m * n_channels
+ downs.append(Downsample(ch_in, with_conv=False))
+
+ self.model = nn.ModuleList(blocks)
+ self.downsampler = nn.ModuleList(downs)
+
+ def instantiate_pretrained(self, config):
+ model = instantiate_from_config(config)
+ self.pretrained_model = model.eval()
+ # self.pretrained_model.train = False
+ for param in self.pretrained_model.parameters():
+ param.requires_grad = False
+
+ @torch.no_grad()
+ def encode_with_pretrained(self, x):
+ c = self.pretrained_model.encode(x)
+ if isinstance(c, DiagonalGaussianDistribution):
+ c = c.mode()
+ return c
+
+ def forward(self, x):
+ z_fs = self.encode_with_pretrained(x)
+ z = self.proj_norm(z_fs)
+ z = self.proj(z)
+ z = nonlinearity(z)
+
+ for submodel, downmodel in zip(self.model, self.downsampler):
+ z = submodel(z, temb=None)
+ z = downmodel(z)
+
+ if self.do_reshape:
+ z = rearrange(z, 'b c h w -> b (h w) c')
+ return z
diff --git a/lidm/modules/diffusion/openaimodel.py b/lidm/modules/diffusion/openaimodel.py
new file mode 100644
index 0000000000000000000000000000000000000000..d21458a9cf84e6e0a797b408d14e3266eeadd704
--- /dev/null
+++ b/lidm/modules/diffusion/openaimodel.py
@@ -0,0 +1,971 @@
+from abc import abstractmethod
+import math
+
+import numpy as np
+import torch as th
+import torch.nn as nn
+import torch.nn.functional as F
+
+from ..basic import (
+ checkpoint,
+ conv_nd,
+ linear,
+ avg_pool_nd,
+ zero_module,
+ normalization,
+ timestep_embedding,
+)
+from ...modules.attention import SpatialTransformer
+
+
+# dummy replace
+def convert_module_to_f16(x):
+ pass
+
+
+def convert_module_to_f32(x):
+ pass
+
+
+## go
+class AttentionPool2d(nn.Module):
+ """
+ Adapted from CLIP: https://github.com/openai/CLIP/blob/main/clip/model.py
+ """
+
+ def __init__(
+ self,
+ spacial_dim: int,
+ embed_dim: int,
+ num_heads_channels: int,
+ output_dim: int = None,
+ ):
+ super().__init__()
+ self.positional_embedding = nn.Parameter(th.randn(embed_dim, spacial_dim ** 2 + 1) / embed_dim ** 0.5)
+ self.qkv_proj = conv_nd(1, embed_dim, 3 * embed_dim, 1)
+ self.c_proj = conv_nd(1, embed_dim, output_dim or embed_dim, 1)
+ self.num_heads = embed_dim // num_heads_channels
+ self.attention = QKVAttention(self.num_heads)
+
+ def forward(self, x):
+ b, c, *_spatial = x.shape
+ x = x.reshape(b, c, -1) # NC(HW)
+ x = th.cat([x.mean(dim=-1, keepdim=True), x], dim=-1) # NC(HW+1)
+ x = x + self.positional_embedding[None, :, :].to(x.dtype) # NC(HW+1)
+ x = self.qkv_proj(x)
+ x = self.attention(x)
+ x = self.c_proj(x)
+ return x[:, :, 0]
+
+
+class TimestepBlock(nn.Module):
+ """
+ Any module where forward() takes timestep embeddings as a second argument.
+ """
+
+ @abstractmethod
+ def forward(self, x, emb):
+ """
+ Apply the module to `x` given `emb` timestep embeddings.
+ """
+
+
+class TimestepEmbedSequential(nn.Sequential, TimestepBlock):
+ """
+ A sequential module that passes timestep embeddings to the children that
+ support it as an extra input.
+ """
+
+ def forward(self, x, emb, context=None):
+ for layer in self:
+ if isinstance(layer, TimestepBlock):
+ x = layer(x, emb)
+ elif isinstance(layer, SpatialTransformer):
+ x = layer(x, context)
+ else:
+ x = layer(x)
+ return x
+
+
+class Upsample(nn.Module):
+ """
+ An upsampling layer with an optional convolution.
+ :param channels: channels in the inputs and outputs.
+ :param use_conv: a bool determining if a convolution is applied.
+ :param dims: determines if the signal is 1D, 2D, or 3D. If 3D, then
+ upsampling occurs in the inner-two dimensions.
+ """
+
+ def __init__(self, channels, use_conv, dims=2, out_channels=None, padding=1, cconv=False):
+ super().__init__()
+ self.channels = channels
+ self.out_channels = out_channels or channels
+ self.use_conv = use_conv
+ self.dims = dims
+ if use_conv:
+ self.conv = conv_nd(dims, self.channels, self.out_channels, 3, padding=padding, cconv=cconv)
+
+ def forward(self, x):
+ assert x.shape[1] == self.channels
+ if self.dims == 3:
+ x = F.interpolate(
+ x, (x.shape[2], x.shape[3] * 2, x.shape[4] * 2), mode="nearest"
+ )
+ else:
+ x = F.interpolate(x, scale_factor=2, mode="nearest")
+ if self.use_conv:
+ x = self.conv(x)
+ return x
+
+
+class TransposedUpsample(nn.Module):
+ 'Learned 2x upsampling without padding'
+
+ def __init__(self, channels, out_channels=None, ks=5):
+ super().__init__()
+ self.channels = channels
+ self.out_channels = out_channels or channels
+
+ self.up = nn.ConvTranspose2d(self.channels, self.out_channels, kernel_size=ks, stride=2)
+
+ def forward(self, x):
+ return self.up(x)
+
+
+class Downsample(nn.Module):
+ """
+ A downsampling layer with an optional convolution.
+ :param channels: channels in the inputs and outputs.
+ :param use_conv: a bool determining if a convolution is applied.
+ :param dims: determines if the signal is 1D, 2D, or 3D. If 3D, then
+ downsampling occurs in the inner-two dimensions.
+ """
+
+ def __init__(self, channels, use_conv, dims=2, out_channels=None, padding=1, cconv=False):
+ super().__init__()
+ self.channels = channels
+ self.out_channels = out_channels or channels
+ self.use_conv = use_conv
+ self.dims = dims
+ stride = 2 if dims != 3 else (1, 2, 2)
+ if use_conv:
+ self.op = conv_nd(
+ dims, self.channels, self.out_channels, 3, stride=stride, padding=padding, cconv=cconv
+ )
+ else:
+ assert self.channels == self.out_channels
+ self.op = avg_pool_nd(dims, kernel_size=stride, stride=stride)
+
+ def forward(self, x):
+ assert x.shape[1] == self.channels
+ return self.op(x)
+
+
+class ResBlock(TimestepBlock):
+ """
+ A residual block that can optionally change the number of channels.
+ :param channels: the number of input channels.
+ :param emb_channels: the number of timestep embedding channels.
+ :param dropout: the rate of dropout.
+ :param out_channels: if specified, the number of out channels.
+ :param use_conv: if True and out_channels is specified, use a spatial
+ convolution instead of a smaller 1x1 convolution to change the
+ channels in the skip connection.
+ :param dims: determines if the signal is 1D, 2D, or 3D.
+ :param use_checkpoint: if True, use gradient checkpointing on this module.
+ :param up: if True, use this block for upsampling.
+ :param down: if True, use this block for downsampling.
+ """
+
+ def __init__(
+ self,
+ channels,
+ emb_channels,
+ dropout,
+ out_channels=None,
+ use_conv=False,
+ use_scale_shift_norm=False,
+ dims=2,
+ use_checkpoint=False,
+ up=False,
+ down=False,
+ cconv=False
+ ):
+ super().__init__()
+ self.channels = channels
+ self.emb_channels = emb_channels
+ self.dropout = dropout
+ self.out_channels = out_channels or channels
+ self.use_conv = use_conv
+ self.use_checkpoint = use_checkpoint
+ self.use_scale_shift_norm = use_scale_shift_norm
+
+ self.in_layers = nn.Sequential(
+ normalization(channels),
+ nn.SiLU(),
+ conv_nd(dims, channels, self.out_channels, 3, padding=1, cconv=cconv),
+ )
+
+ self.updown = up or down
+
+ if up:
+ self.h_upd = Upsample(channels, False, dims, cconv=cconv)
+ self.x_upd = Upsample(channels, False, dims, cconv=cconv)
+ elif down:
+ self.h_upd = Downsample(channels, False, dims, cconv=cconv)
+ self.x_upd = Downsample(channels, False, dims, cconv=cconv)
+ else:
+ self.h_upd = self.x_upd = nn.Identity()
+
+ self.emb_layers = nn.Sequential(
+ nn.SiLU(),
+ linear(
+ emb_channels,
+ 2 * self.out_channels if use_scale_shift_norm else self.out_channels,
+ ),
+ )
+ self.out_layers = nn.Sequential(
+ normalization(self.out_channels),
+ nn.SiLU(),
+ nn.Dropout(p=dropout),
+ zero_module(
+ conv_nd(dims, self.out_channels, self.out_channels, 3, padding=1, cconv=cconv)
+ ),
+ )
+
+ if self.out_channels == channels:
+ self.skip_connection = nn.Identity()
+ elif use_conv:
+ self.skip_connection = conv_nd(
+ dims, channels, self.out_channels, 3, padding=1, cconv=cconv
+ )
+ else:
+ self.skip_connection = conv_nd(dims, channels, self.out_channels, 1, cconv=cconv)
+
+ def forward(self, x, emb):
+ """
+ Apply the block to a Tensor, conditioned on a timestep embedding.
+ :param x: an [N x C x ...] Tensor of features.
+ :param emb: an [N x emb_channels] Tensor of timestep embeddings.
+ :return: an [N x C x ...] Tensor of outputs.
+ """
+ return checkpoint(
+ self._forward, (x, emb), self.parameters(), self.use_checkpoint
+ )
+
+ def _forward(self, x, emb):
+ if self.updown:
+ in_rest, in_conv = self.in_layers[:-1], self.in_layers[-1]
+ h = in_rest(x)
+ h = self.h_upd(h)
+ x = self.x_upd(x)
+ h = in_conv(h)
+ else:
+ h = self.in_layers(x)
+ emb_out = self.emb_layers(emb).type(h.dtype)
+ while len(emb_out.shape) < len(h.shape):
+ emb_out = emb_out[..., None]
+ if self.use_scale_shift_norm:
+ out_norm, out_rest = self.out_layers[0], self.out_layers[1:]
+ scale, shift = th.chunk(emb_out, 2, dim=1)
+ h = out_norm(h) * (1 + scale) + shift
+ h = out_rest(h)
+ else:
+ h = h + emb_out
+ h = self.out_layers(h)
+ return self.skip_connection(x) + h
+
+
+class AttentionBlock(nn.Module):
+ """
+ An attention block that allows spatial positions to attend to each other.
+ Originally ported from here, but adapted to the N-d case.
+ https://github.com/hojonathanho/diffusion/blob/1e0dceb3b3495bbe19116a5e1b3596cd0706c543/diffusion_tf/models/unet.py#L66.
+ """
+
+ def __init__(
+ self,
+ channels,
+ num_heads=1,
+ num_head_channels=-1,
+ use_checkpoint=False,
+ use_new_attention_order=False,
+ ):
+ super().__init__()
+ self.channels = channels
+ if num_head_channels == -1:
+ self.num_heads = num_heads
+ else:
+ assert (
+ channels % num_head_channels == 0
+ ), f"q,k,v channels {channels} is not divisible by num_head_channels {num_head_channels}"
+ self.num_heads = channels // num_head_channels
+ self.use_checkpoint = use_checkpoint
+ self.norm = normalization(channels)
+ self.qkv = conv_nd(1, channels, channels * 3, 1)
+ if use_new_attention_order:
+ # split qkv before split heads
+ self.attention = QKVAttention(self.num_heads)
+ else:
+ # split heads before split qkv
+ self.attention = QKVAttentionLegacy(self.num_heads)
+
+ self.proj_out = zero_module(conv_nd(1, channels, channels, 1))
+
+ def forward(self, x):
+ return checkpoint(self._forward, (x,), self.parameters(),
+ True) # TODO: check checkpoint usage, is True # TODO: fix the .half call!!!
+ # return pt_checkpoint(self._forward, x) # pytorch
+
+ def _forward(self, x):
+ b, c, *spatial = x.shape
+ x = x.reshape(b, c, -1)
+ qkv = self.qkv(self.norm(x))
+ h = self.attention(qkv)
+ h = self.proj_out(h)
+ return (x + h).reshape(b, c, *spatial)
+
+
+def count_flops_attn(model, _x, y):
+ """
+ A counter for the `thop` package to count the operations in an
+ attention operation.
+ Meant to be used like:
+ macs, params = thop.profile(
+ model,
+ inputs=(inputs, timestamps),
+ custom_ops={QKVAttention: QKVAttention.count_flops},
+ )
+ """
+ b, c, *spatial = y[0].shape
+ num_spatial = int(np.prod(spatial))
+ # We perform two matmuls with the same number of ops.
+ # The first computes the weight matrix, the second computes
+ # the combination of the value vectors.
+ matmul_ops = 2 * b * (num_spatial ** 2) * c
+ model.total_ops += th.DoubleTensor([matmul_ops])
+
+
+class QKVAttentionLegacy(nn.Module):
+ """
+ A module which performs QKV attention. Matches legacy QKVAttention + input/ouput heads shaping
+ """
+
+ def __init__(self, n_heads):
+ super().__init__()
+ self.n_heads = n_heads
+
+ def forward(self, qkv):
+ """
+ Apply QKV attention.
+ :param qkv: an [N x (H * 3 * C) x T] tensor of Qs, Ks, and Vs.
+ :return: an [N x (H * C) x T] tensor after attention.
+ """
+ bs, width, length = qkv.shape
+ assert width % (3 * self.n_heads) == 0
+ ch = width // (3 * self.n_heads)
+ q, k, v = qkv.reshape(bs * self.n_heads, ch * 3, length).split(ch, dim=1)
+ scale = 1 / math.sqrt(math.sqrt(ch))
+ weight = th.einsum(
+ "bct,bcs->bts", q * scale, k * scale
+ ) # More stable with f16 than dividing afterwards
+ weight = th.softmax(weight.float(), dim=-1).type(weight.dtype)
+ a = th.einsum("bts,bcs->bct", weight, v)
+ return a.reshape(bs, -1, length)
+
+ @staticmethod
+ def count_flops(model, _x, y):
+ return count_flops_attn(model, _x, y)
+
+
+class QKVAttention(nn.Module):
+ """
+ A module which performs QKV attention and splits in a different order.
+ """
+
+ def __init__(self, n_heads):
+ super().__init__()
+ self.n_heads = n_heads
+
+ def forward(self, qkv):
+ """
+ Apply QKV attention.
+ :param qkv: an [N x (3 * H * C) x T] tensor of Qs, Ks, and Vs.
+ :return: an [N x (H * C) x T] tensor after attention.
+ """
+ bs, width, length = qkv.shape
+ assert width % (3 * self.n_heads) == 0
+ ch = width // (3 * self.n_heads)
+ q, k, v = qkv.chunk(3, dim=1)
+ scale = 1 / math.sqrt(math.sqrt(ch))
+ weight = th.einsum(
+ "bct,bcs->bts",
+ (q * scale).view(bs * self.n_heads, ch, length),
+ (k * scale).view(bs * self.n_heads, ch, length),
+ ) # More stable with f16 than dividing afterwards
+ weight = th.softmax(weight.float(), dim=-1).type(weight.dtype)
+ a = th.einsum("bts,bcs->bct", weight, v.reshape(bs * self.n_heads, ch, length))
+ return a.reshape(bs, -1, length)
+
+ @staticmethod
+ def count_flops(model, _x, y):
+ return count_flops_attn(model, _x, y)
+
+
+class UNetModel(nn.Module):
+ """
+ The full UNet model with attention and timestep embedding.
+ :param in_channels: channels in the input Tensor.
+ :param model_channels: base channel count for the model.
+ :param out_channels: channels in the output Tensor.
+ :param num_res_blocks: number of residual blocks per downsample.
+ :param attention_resolutions: a collection of downsample rates at which
+ attention will take place. May be a set, list, or tuple.
+ For example, if this contains 4, then at 4x downsampling, attention
+ will be used.
+ :param dropout: the dropout probability.
+ :param channel_mult: channel multiplier for each level of the UNet.
+ :param conv_resample: if True, use learned convolutions for upsampling and
+ downsampling.
+ :param dims: determines if the signal is 1D, 2D, or 3D.
+ :param num_classes: if specified (as an int), then this model will be
+ class-conditional with `num_classes` classes.
+ :param use_checkpoint: use gradient checkpointing to reduce memory usage.
+ :param num_heads: the number of attention heads in each attention layer.
+ :param num_heads_channels: if specified, ignore num_heads and instead use
+ a fixed channel width per attention head.
+ :param num_heads_upsample: works with num_heads to set a different number
+ of heads for upsampling. Deprecated.
+ :param use_scale_shift_norm: use a FiLM-like conditioning mechanism.
+ :param resblock_updown: use residual blocks for up/downsampling.
+ :param use_new_attention_order: use a different attention pattern for potentially
+ increased efficiency.
+ """
+
+ def __init__(
+ self,
+ image_size,
+ in_channels,
+ model_channels,
+ out_channels,
+ num_res_blocks,
+ attention_resolutions,
+ dropout=0,
+ channel_mult=(1, 2, 4, 8),
+ conv_resample=True,
+ dims=2,
+ num_classes=None,
+ use_checkpoint=False,
+ use_fp16=False,
+ num_heads=-1,
+ num_head_channels=-1,
+ num_heads_upsample=-1,
+ use_scale_shift_norm=False,
+ resblock_updown=False,
+ use_new_attention_order=False,
+ use_spatial_transformer=False, # custom transformer support
+ transformer_depth=1, # custom transformer support
+ context_dim=None, # custom transformer support
+ n_embed=None, # custom support for prediction of discrete ids into codebook of first stage vq model
+ legacy=True,
+ lib_name='ldm'
+ ):
+ super().__init__()
+ if use_spatial_transformer:
+ assert context_dim is not None, 'Fool!! You forgot to include the dimension of your cross-attention conditioning...'
+ if context_dim is not None:
+ assert use_spatial_transformer, 'Fool!! You forgot to use the spatial transformer for your cross-attention conditioning...'
+ from omegaconf.listconfig import ListConfig
+ if type(context_dim) == ListConfig:
+ context_dim = list(context_dim)
+
+ if num_heads_upsample == -1:
+ num_heads_upsample = num_heads
+
+ if num_heads == -1:
+ assert num_head_channels != -1, 'Either num_heads or num_head_channels has to be set'
+
+ if num_head_channels == -1:
+ assert num_heads != -1, 'Either num_heads or num_head_channels has to be set'
+
+ self.image_size = image_size
+ self.in_channels = in_channels
+ self.model_channels = model_channels
+ self.out_channels = out_channels
+ self.num_res_blocks = num_res_blocks
+ self.attention_resolutions = attention_resolutions
+ self.dropout = dropout
+ self.channel_mult = channel_mult
+ self.conv_resample = conv_resample
+ self.num_classes = num_classes
+ self.use_checkpoint = use_checkpoint
+ self.dtype = th.float16 if use_fp16 else th.float32
+ self.num_heads = num_heads
+ self.num_head_channels = num_head_channels
+ self.num_heads_upsample = num_heads_upsample
+ self.predict_codebook_ids = n_embed is not None
+ self.cconv = lib_name in ['lidm', 'lidm_v0']
+
+ time_embed_dim = model_channels * 4
+ self.time_embed = nn.Sequential(
+ linear(model_channels, time_embed_dim),
+ nn.SiLU(),
+ linear(time_embed_dim, time_embed_dim),
+ )
+
+ if self.num_classes is not None:
+ self.label_emb = nn.Embedding(num_classes, time_embed_dim)
+
+ self.input_blocks = nn.ModuleList(
+ [
+ TimestepEmbedSequential(
+ conv_nd(dims, in_channels, model_channels, 3, padding=1, cconv=self.cconv)
+ )
+ ]
+ )
+ self._feature_size = model_channels
+ input_block_chans = [model_channels]
+ ch = model_channels
+ ds = 1
+ for level, mult in enumerate(channel_mult):
+ for _ in range(num_res_blocks):
+ layers = [
+ ResBlock(
+ ch,
+ time_embed_dim,
+ dropout,
+ out_channels=mult * model_channels,
+ dims=dims,
+ use_checkpoint=use_checkpoint,
+ use_scale_shift_norm=use_scale_shift_norm,
+ cconv=self.cconv
+ )
+ ]
+ ch = mult * model_channels
+ if ds in attention_resolutions:
+ if num_head_channels == -1:
+ dim_head = ch // num_heads
+ else:
+ num_heads = ch // num_head_channels
+ dim_head = num_head_channels
+ if legacy:
+ # num_heads = 1
+ dim_head = ch // num_heads if use_spatial_transformer else num_head_channels
+ layers.append(
+ AttentionBlock(
+ ch,
+ use_checkpoint=use_checkpoint,
+ num_heads=num_heads,
+ num_head_channels=dim_head,
+ use_new_attention_order=use_new_attention_order,
+ ) if not use_spatial_transformer else SpatialTransformer(
+ ch, num_heads, dim_head, depth=transformer_depth, context_dim=context_dim
+ )
+ )
+ self.input_blocks.append(TimestepEmbedSequential(*layers))
+ self._feature_size += ch
+ input_block_chans.append(ch)
+ if level != len(channel_mult) - 1:
+ out_ch = ch
+ self.input_blocks.append(
+ TimestepEmbedSequential(
+ ResBlock(
+ ch,
+ time_embed_dim,
+ dropout,
+ out_channels=out_ch,
+ dims=dims,
+ use_checkpoint=use_checkpoint,
+ use_scale_shift_norm=use_scale_shift_norm,
+ down=True,
+ cconv=self.cconv
+ )
+ if resblock_updown
+ else Downsample(
+ ch, conv_resample, dims=dims, out_channels=out_ch, cconv=self.cconv
+ )
+ )
+ )
+ ch = out_ch
+ input_block_chans.append(ch)
+ ds *= 2
+ self._feature_size += ch
+
+ if num_head_channels == -1:
+ dim_head = ch // num_heads
+ else:
+ num_heads = ch // num_head_channels
+ dim_head = num_head_channels
+ if legacy:
+ # num_heads = 1
+ dim_head = ch // num_heads if use_spatial_transformer else num_head_channels
+ self.middle_block = TimestepEmbedSequential(
+ ResBlock(
+ ch,
+ time_embed_dim,
+ dropout,
+ dims=dims,
+ use_checkpoint=use_checkpoint,
+ use_scale_shift_norm=use_scale_shift_norm,
+ cconv=self.cconv
+ ),
+ AttentionBlock(
+ ch,
+ use_checkpoint=use_checkpoint,
+ num_heads=num_heads,
+ num_head_channels=dim_head,
+ use_new_attention_order=use_new_attention_order,
+ ) if not use_spatial_transformer else SpatialTransformer(
+ ch, num_heads, dim_head, depth=transformer_depth, context_dim=context_dim
+ ),
+ ResBlock(
+ ch,
+ time_embed_dim,
+ dropout,
+ dims=dims,
+ use_checkpoint=use_checkpoint,
+ use_scale_shift_norm=use_scale_shift_norm,
+ cconv=self.cconv
+ ),
+ )
+ self._feature_size += ch
+
+ self.output_blocks = nn.ModuleList([])
+ for level, mult in list(enumerate(channel_mult))[::-1]:
+ for i in range(num_res_blocks + 1):
+ ich = input_block_chans.pop()
+ layers = [
+ ResBlock(
+ ch + ich,
+ time_embed_dim,
+ dropout,
+ out_channels=model_channels * mult,
+ dims=dims,
+ use_checkpoint=use_checkpoint,
+ use_scale_shift_norm=use_scale_shift_norm,
+ cconv=self.cconv
+ )
+ ]
+ ch = model_channels * mult
+ if ds in attention_resolutions:
+ if num_head_channels == -1:
+ dim_head = ch // num_heads
+ else:
+ num_heads = ch // num_head_channels
+ dim_head = num_head_channels
+ if legacy:
+ # num_heads = 1
+ dim_head = ch // num_heads if use_spatial_transformer else num_head_channels
+ layers.append(
+ AttentionBlock(
+ ch,
+ use_checkpoint=use_checkpoint,
+ num_heads=num_heads_upsample,
+ num_head_channels=dim_head,
+ use_new_attention_order=use_new_attention_order,
+ ) if not use_spatial_transformer else SpatialTransformer(
+ ch, num_heads, dim_head, depth=transformer_depth, context_dim=context_dim
+ )
+ )
+ if level and i == num_res_blocks:
+ out_ch = ch
+ layers.append(
+ ResBlock(
+ ch,
+ time_embed_dim,
+ dropout,
+ out_channels=out_ch,
+ dims=dims,
+ use_checkpoint=use_checkpoint,
+ use_scale_shift_norm=use_scale_shift_norm,
+ up=True,
+ cconv=self.cconv
+ )
+ if resblock_updown
+ else Upsample(ch, conv_resample, dims=dims, out_channels=out_ch, cconv=self.cconv)
+ )
+ ds //= 2
+ self.output_blocks.append(TimestepEmbedSequential(*layers))
+ self._feature_size += ch
+
+ self.out = nn.Sequential(
+ normalization(ch),
+ nn.SiLU(),
+ zero_module(conv_nd(dims, model_channels, out_channels, 3, padding=1, cconv=self.cconv)),
+ )
+ if self.predict_codebook_ids:
+ self.id_predictor = nn.Sequential(
+ normalization(ch),
+ conv_nd(dims, model_channels, n_embed, 1, use_cconv=self.use_cconv),
+ # nn.LogSoftmax(dim=1) # change to cross_entropy and produce non-normalized logits
+ )
+
+ def convert_to_fp16(self):
+ """
+ Convert the torso of the model to float16.
+ """
+ self.input_blocks.apply(convert_module_to_f16)
+ self.middle_block.apply(convert_module_to_f16)
+ self.output_blocks.apply(convert_module_to_f16)
+
+ def convert_to_fp32(self):
+ """
+ Convert the torso of the model to float32.
+ """
+ self.input_blocks.apply(convert_module_to_f32)
+ self.middle_block.apply(convert_module_to_f32)
+ self.output_blocks.apply(convert_module_to_f32)
+
+ def forward(self, x, timesteps=None, context=None, y=None, **kwargs):
+ """
+ Apply the model to an input batch.
+ :param x: an [N x C x ...] Tensor of inputs.
+ :param timesteps: a 1-D batch of timesteps.
+ :param context: conditioning plugged in via crossattn
+ :param y: an [N] Tensor of labels, if class-conditional.
+ :return: an [N x C x ...] Tensor of outputs.
+ """
+ assert (y is not None) == (
+ self.num_classes is not None
+ ), "must specify y if and only if the model is class-conditional"
+ hs = []
+ t_emb = timestep_embedding(timesteps, self.model_channels, repeat_only=False)
+ emb = self.time_embed(t_emb)
+
+ if self.num_classes is not None:
+ assert y.shape == (x.shape[0],)
+ emb = emb + self.label_emb(y)
+
+ h = x.type(self.dtype)
+ for module in self.input_blocks:
+ h = module(h, emb, context)
+ hs.append(h)
+ h = self.middle_block(h, emb, context)
+ for module in self.output_blocks:
+ h = th.cat([h, hs.pop()], dim=1)
+ h = module(h, emb, context)
+ h = h.type(x.dtype)
+ if self.predict_codebook_ids:
+ return self.id_predictor(h)
+ else:
+ return self.out(h)
+
+
+class EncoderUNetModel(nn.Module):
+ """
+ The half UNet model with attention and timestep embedding.
+ For usage, see UNet.
+ """
+
+ def __init__(
+ self,
+ image_size,
+ in_channels,
+ model_channels,
+ out_channels,
+ num_res_blocks,
+ attention_resolutions,
+ dropout=0,
+ channel_mult=(1, 2, 4, 8),
+ conv_resample=True,
+ dims=2,
+ use_checkpoint=False,
+ use_fp16=False,
+ num_heads=1,
+ num_head_channels=-1,
+ num_heads_upsample=-1,
+ use_scale_shift_norm=False,
+ resblock_updown=False,
+ use_new_attention_order=False,
+ pool="adaptive",
+ lib_name='ldm',
+ *args,
+ **kwargs
+ ):
+ super().__init__()
+
+ if num_heads_upsample == -1:
+ num_heads_upsample = num_heads
+
+ self.in_channels = in_channels
+ self.model_channels = model_channels
+ self.out_channels = out_channels
+ self.num_res_blocks = num_res_blocks
+ self.attention_resolutions = attention_resolutions
+ self.dropout = dropout
+ self.channel_mult = channel_mult
+ self.conv_resample = conv_resample
+ self.use_checkpoint = use_checkpoint
+ self.dtype = th.float16 if use_fp16 else th.float32
+ self.num_heads = num_heads
+ self.num_head_channels = num_head_channels
+ self.num_heads_upsample = num_heads_upsample
+ self.cconv = lib_name == 'lidm'
+
+ time_embed_dim = model_channels * 4
+ self.time_embed = nn.Sequential(
+ linear(model_channels, time_embed_dim),
+ nn.SiLU(),
+ linear(time_embed_dim, time_embed_dim),
+ )
+
+ self.input_blocks = nn.ModuleList(
+ [
+ TimestepEmbedSequential(
+ conv_nd(dims, in_channels, model_channels, 3, padding=1, cconv=self.cconv)
+ )
+ ]
+ )
+ self._feature_size = model_channels
+ input_block_chans = [model_channels]
+ ch = model_channels
+ ds = 1
+ for level, mult in enumerate(channel_mult):
+ for _ in range(num_res_blocks):
+ layers = [
+ ResBlock(
+ ch,
+ time_embed_dim,
+ dropout,
+ out_channels=mult * model_channels,
+ dims=dims,
+ use_checkpoint=use_checkpoint,
+ use_scale_shift_norm=use_scale_shift_norm,
+ )
+ ]
+ ch = mult * model_channels
+ if ds in attention_resolutions:
+ layers.append(
+ AttentionBlock(
+ ch,
+ use_checkpoint=use_checkpoint,
+ num_heads=num_heads,
+ num_head_channels=num_head_channels,
+ use_new_attention_order=use_new_attention_order,
+ )
+ )
+ self.input_blocks.append(TimestepEmbedSequential(*layers))
+ self._feature_size += ch
+ input_block_chans.append(ch)
+ if level != len(channel_mult) - 1:
+ out_ch = ch
+ self.input_blocks.append(
+ TimestepEmbedSequential(
+ ResBlock(
+ ch,
+ time_embed_dim,
+ dropout,
+ out_channels=out_ch,
+ dims=dims,
+ use_checkpoint=use_checkpoint,
+ use_scale_shift_norm=use_scale_shift_norm,
+ down=True,
+ )
+ if resblock_updown
+ else Downsample(
+ ch, conv_resample, dims=dims, out_channels=out_ch
+ )
+ )
+ )
+ ch = out_ch
+ input_block_chans.append(ch)
+ ds *= 2
+ self._feature_size += ch
+
+ self.middle_block = TimestepEmbedSequential(
+ ResBlock(
+ ch,
+ time_embed_dim,
+ dropout,
+ dims=dims,
+ use_checkpoint=use_checkpoint,
+ use_scale_shift_norm=use_scale_shift_norm,
+ ),
+ AttentionBlock(
+ ch,
+ use_checkpoint=use_checkpoint,
+ num_heads=num_heads,
+ num_head_channels=num_head_channels,
+ use_new_attention_order=use_new_attention_order,
+ ),
+ ResBlock(
+ ch,
+ time_embed_dim,
+ dropout,
+ dims=dims,
+ use_checkpoint=use_checkpoint,
+ use_scale_shift_norm=use_scale_shift_norm,
+ ),
+ )
+ self._feature_size += ch
+ self.pool = pool
+ if pool == "adaptive":
+ self.out = nn.Sequential(
+ normalization(ch),
+ nn.SiLU(),
+ nn.AdaptiveAvgPool2d((1, 1)),
+ zero_module(conv_nd(dims, ch, out_channels, 1)),
+ nn.Flatten(),
+ )
+ elif pool == "attention":
+ assert num_head_channels != -1
+ self.out = nn.Sequential(
+ normalization(ch),
+ nn.SiLU(),
+ AttentionPool2d(
+ (image_size // ds), ch, num_head_channels, out_channels
+ ),
+ )
+ elif pool == "spatial":
+ self.out = nn.Sequential(
+ nn.Linear(self._feature_size, 2048),
+ nn.ReLU(),
+ nn.Linear(2048, self.out_channels),
+ )
+ elif pool == "spatial_v2":
+ self.out = nn.Sequential(
+ nn.Linear(self._feature_size, 2048),
+ normalization(2048),
+ nn.SiLU(),
+ nn.Linear(2048, self.out_channels),
+ )
+ else:
+ raise NotImplementedError(f"Unexpected {pool} pooling")
+
+ def convert_to_fp16(self):
+ """
+ Convert the torso of the model to float16.
+ """
+ self.input_blocks.apply(convert_module_to_f16)
+ self.middle_block.apply(convert_module_to_f16)
+
+ def convert_to_fp32(self):
+ """
+ Convert the torso of the model to float32.
+ """
+ self.input_blocks.apply(convert_module_to_f32)
+ self.middle_block.apply(convert_module_to_f32)
+
+ def forward(self, x, timesteps):
+ """
+ Apply the model to an input batch.
+ :param x: an [N x C x ...] Tensor of inputs.
+ :param timesteps: a 1-D batch of timesteps.
+ :return: an [N x K] Tensor of outputs.
+ """
+ emb = self.time_embed(timestep_embedding(timesteps, self.model_channels))
+
+ results = []
+ h = x.type(self.dtype)
+ for module in self.input_blocks:
+ h = module(h, emb)
+ if self.pool.startswith("spatial"):
+ results.append(h.type(x.dtype).mean(dim=(2, 3)))
+ h = self.middle_block(h, emb)
+ if self.pool.startswith("spatial"):
+ results.append(h.type(x.dtype).mean(dim=(2, 3)))
+ h = th.cat(results, axis=-1)
+ return self.out(h)
+ else:
+ h = h.type(x.dtype)
+ return self.out(h)
diff --git a/lidm/modules/distributions/__init__.py b/lidm/modules/distributions/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/lidm/modules/distributions/distributions.py b/lidm/modules/distributions/distributions.py
new file mode 100644
index 0000000000000000000000000000000000000000..f2b8ef901130efc171aa69742ca0244d94d3f2e9
--- /dev/null
+++ b/lidm/modules/distributions/distributions.py
@@ -0,0 +1,92 @@
+import torch
+import numpy as np
+
+
+class AbstractDistribution:
+ def sample(self):
+ raise NotImplementedError()
+
+ def mode(self):
+ raise NotImplementedError()
+
+
+class DiracDistribution(AbstractDistribution):
+ def __init__(self, value):
+ self.value = value
+
+ def sample(self):
+ return self.value
+
+ def mode(self):
+ return self.value
+
+
+class DiagonalGaussianDistribution(object):
+ def __init__(self, parameters, deterministic=False):
+ self.parameters = parameters
+ self.mean, self.logvar = torch.chunk(parameters, 2, dim=1)
+ self.logvar = torch.clamp(self.logvar, -30.0, 20.0)
+ self.deterministic = deterministic
+ self.std = torch.exp(0.5 * self.logvar)
+ self.var = torch.exp(self.logvar)
+ if self.deterministic:
+ self.var = self.std = torch.zeros_like(self.mean).to(device=self.parameters.device)
+
+ def sample(self):
+ x = self.mean + self.std * torch.randn(self.mean.shape).to(device=self.parameters.device)
+ return x
+
+ def kl(self, other=None):
+ if self.deterministic:
+ return torch.Tensor([0.])
+ else:
+ if other is None:
+ return 0.5 * torch.sum(torch.pow(self.mean, 2)
+ + self.var - 1.0 - self.logvar,
+ dim=[1, 2, 3])
+ else:
+ return 0.5 * torch.sum(
+ torch.pow(self.mean - other.mean, 2) / other.var
+ + self.var / other.var - 1.0 - self.logvar + other.logvar,
+ dim=[1, 2, 3])
+
+ def nll(self, sample, dims=[1,2,3]):
+ if self.deterministic:
+ return torch.Tensor([0.])
+ logtwopi = np.log(2.0 * np.pi)
+ return 0.5 * torch.sum(
+ logtwopi + self.logvar + torch.pow(sample - self.mean, 2) / self.var,
+ dim=dims)
+
+ def mode(self):
+ return self.mean
+
+
+def normal_kl(mean1, logvar1, mean2, logvar2):
+ """
+ source: https://github.com/openai/guided-diffusion/blob/27c20a8fab9cb472df5d6bdd6c8d11c8f430b924/guided_diffusion/losses.py#L12
+ Compute the KL divergence between two gaussians.
+ Shapes are automatically broadcasted, so batches can be compared to
+ scalars, among other use cases.
+ """
+ tensor = None
+ for obj in (mean1, logvar1, mean2, logvar2):
+ if isinstance(obj, torch.Tensor):
+ tensor = obj
+ break
+ assert tensor is not None, "at least one argument must be a Tensor"
+
+ # Force variances to be Tensors. Broadcasting helps convert scalars to
+ # Tensors, but it does not work for torch.exp().
+ logvar1, logvar2 = [
+ x if isinstance(x, torch.Tensor) else torch.tensor(x).to(tensor)
+ for x in (logvar1, logvar2)
+ ]
+
+ return 0.5 * (
+ -1.0
+ + logvar2
+ - logvar1
+ + torch.exp(logvar1 - logvar2)
+ + ((mean1 - mean2) ** 2) * torch.exp(-logvar2)
+ )
diff --git a/lidm/modules/ema.py b/lidm/modules/ema.py
new file mode 100644
index 0000000000000000000000000000000000000000..ea6d1d3594bb6336ed896557cf33275ea94380d3
--- /dev/null
+++ b/lidm/modules/ema.py
@@ -0,0 +1,76 @@
+import torch
+from torch import nn
+
+
+class LitEma(nn.Module):
+ def __init__(self, model, decay=0.9999, use_num_upates=True):
+ super().__init__()
+ if decay < 0.0 or decay > 1.0:
+ raise ValueError('Decay must be between 0 and 1')
+
+ self.m_name2s_name = {}
+ self.register_buffer('decay', torch.tensor(decay, dtype=torch.float32))
+ self.register_buffer('num_updates', torch.tensor(0, dtype=torch.int) if use_num_upates
+ else torch.tensor(-1, dtype=torch.int))
+
+ for name, p in model.named_parameters():
+ if p.requires_grad:
+ # remove as '.'-character is not allowed in buffers
+ s_name = name.replace('.', '')
+ self.m_name2s_name.update({name: s_name})
+ self.register_buffer(s_name, p.clone().detach().data)
+
+ self.collected_params = []
+
+ def forward(self, model):
+ decay = self.decay
+
+ if self.num_updates >= 0:
+ self.num_updates += 1
+ decay = min(self.decay, (1 + self.num_updates) / (10 + self.num_updates))
+
+ one_minus_decay = 1.0 - decay
+
+ with torch.no_grad():
+ m_param = dict(model.named_parameters())
+ shadow_params = dict(self.named_buffers())
+
+ for key in m_param:
+ if m_param[key].requires_grad:
+ sname = self.m_name2s_name[key]
+ shadow_params[sname] = shadow_params[sname].type_as(m_param[key])
+ shadow_params[sname].sub_(one_minus_decay * (shadow_params[sname] - m_param[key]))
+ else:
+ assert not key in self.m_name2s_name
+
+ def copy_to(self, model):
+ m_param = dict(model.named_parameters())
+ shadow_params = dict(self.named_buffers())
+ for key in m_param:
+ if m_param[key].requires_grad:
+ m_param[key].data.copy_(shadow_params[self.m_name2s_name[key]].data)
+ else:
+ assert not key in self.m_name2s_name
+
+ def store(self, parameters):
+ """
+ Save the current parameters for restoring later.
+ Args:
+ parameters: Iterable of `torch.nn.Parameter`; the parameters to be
+ temporarily stored.
+ """
+ self.collected_params = [param.clone() for param in parameters]
+
+ def restore(self, parameters):
+ """
+ Restore the parameters stored with the `store` method.
+ Useful to validate the model with EMA parameters without affecting the
+ original optimization process. Store the parameters before the
+ `copy_to` method. After validation (or model saving), use this to
+ restore the former parameters.
+ Args:
+ parameters: Iterable of `torch.nn.Parameter`; the parameters to be
+ updated with the stored parameters.
+ """
+ for c_param, param in zip(self.collected_params, parameters):
+ param.data.copy_(c_param.data)
diff --git a/lidm/modules/encoders/__init__.py b/lidm/modules/encoders/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/lidm/modules/encoders/modules.py b/lidm/modules/encoders/modules.py
new file mode 100644
index 0000000000000000000000000000000000000000..fb37ae10aec342c6ff03fc1ecc979588dd92aa0a
--- /dev/null
+++ b/lidm/modules/encoders/modules.py
@@ -0,0 +1,327 @@
+import torch
+import torch.nn as nn
+from functools import partial
+import clip
+from einops import rearrange, repeat
+import kornia
+
+from ...modules.x_transformer import Encoder, TransformerWrapper
+
+
+class AbstractEncoder(nn.Module):
+ def __init__(self):
+ super().__init__()
+
+ def encode(self, *args, **kwargs):
+ raise NotImplementedError
+
+
+class ClassEmbedder(nn.Module):
+ def __init__(self, embed_dim, n_classes=1000, key='class'):
+ super().__init__()
+ self.key = key
+ self.embedding = nn.Embedding(n_classes, embed_dim)
+
+ def forward(self, batch, key=None):
+ if key is None:
+ key = self.key
+ # this is for use in crossattn
+ c = batch[key][:, None]
+ c = self.embedding(c)
+ return c
+
+
+class TransformerEmbedder(AbstractEncoder):
+ """Some transformer encoder layers"""
+
+ def __init__(self, n_embed, n_layer, vocab_size, max_seq_len=77, device="cuda"):
+ super().__init__()
+ self.device = device
+ self.transformer = TransformerWrapper(num_tokens=vocab_size, max_seq_len=max_seq_len,
+ attn_layers=Encoder(dim=n_embed, depth=n_layer))
+
+ def forward(self, tokens):
+ tokens = tokens.to(self.device) # meh
+ z = self.transformer(tokens, return_embeddings=True)
+ return z
+
+ def encode(self, x):
+ return self(x)
+
+
+class BERTTokenizer(AbstractEncoder):
+ """ Uses a pretrained BERT tokenizer by huggingface. Vocab size: 30522 (?)"""
+
+ def __init__(self, device="cuda", vq_interface=True, max_length=77):
+ super().__init__()
+ from transformers import BertTokenizerFast
+ # self.tokenizer = BertTokenizerFast.from_pretrained("bert-base-uncased")
+ self.tokenizer = BertTokenizerFast.from_pretrained('./models/bert')
+ self.device = device
+ self.vq_interface = vq_interface
+ self.max_length = max_length
+
+ def forward(self, text):
+ batch_encoding = self.tokenizer(text, truncation=True, max_length=self.max_length, return_length=True,
+ return_overflowing_tokens=False, padding="max_length", return_tensors="pt")
+ tokens = batch_encoding["input_ids"].to(self.device)
+ return tokens
+
+ @torch.no_grad()
+ def encode(self, text):
+ tokens = self(text)
+ if not self.vq_interface:
+ return tokens
+ return None, None, [None, None, tokens]
+
+ def decode(self, text):
+ return text
+
+
+class BERTEmbedder(AbstractEncoder):
+ """Uses the BERT tokenizr model and add some transformer encoder layers"""
+
+ def __init__(self, n_embed, n_layer, vocab_size=30522, max_seq_len=77,
+ device="cuda", use_tokenizer=True, embedding_dropout=0.0):
+ super().__init__()
+ self.use_tknz_fn = use_tokenizer
+ if self.use_tknz_fn:
+ self.tknz_fn = BERTTokenizer(vq_interface=False, max_length=max_seq_len)
+ self.device = device
+ self.transformer = TransformerWrapper(num_tokens=vocab_size, max_seq_len=max_seq_len,
+ attn_layers=Encoder(dim=n_embed, depth=n_layer),
+ emb_dropout=embedding_dropout)
+
+ def forward(self, text):
+ if self.use_tknz_fn:
+ tokens = self.tknz_fn(text) # .to(self.device)
+ else:
+ tokens = text
+ z = self.transformer(tokens, return_embeddings=True)
+ return z
+
+ def encode(self, text):
+ # output of length 77
+ return self(text)
+
+
+class SpatialRescaler(nn.Module):
+ def __init__(self,
+ strides=[],
+ method='bilinear',
+ in_channels=3,
+ out_channels=None,
+ bias=False):
+ super().__init__()
+ self.strides = strides
+ assert method in ['nearest', 'linear', 'bilinear', 'trilinear', 'bicubic', 'area']
+ self.interpolator = partial(torch.nn.functional.interpolate, mode=method, align_corners=True)
+ self.remap_output = out_channels is not None
+ if self.remap_output:
+ print(f'Spatial Rescaler mapping from {in_channels} to {out_channels} channels after resizing.')
+ self.channel_mapper = nn.Conv2d(in_channels, out_channels, 1, bias=bias)
+
+ def forward(self, x):
+ for h_s, w_s in self.strides:
+ x = self.interpolator(x, scale_factor=(1/h_s, 1/w_s))
+
+ if self.remap_output:
+ x = self.channel_mapper(x)
+ return x
+
+ def encode(self, x):
+ return self(x)
+
+
+class FrozenCLIPTextEmbedder(nn.Module):
+ """
+ Uses the CLIP transformer encoder for text.
+ """
+
+ def __init__(self, version='ViT-L/14', device="cuda", max_length=77, n_repeat=1, normalize=True):
+ super().__init__()
+ self.model, _ = clip.load(version, jit=False, device="cpu")
+ self.model.to(device)
+ self.device = device
+ self.max_length = max_length
+ self.n_repeat = n_repeat
+ self.normalize = normalize
+
+ def freeze(self):
+ self.model = self.model.eval()
+ for param in self.parameters():
+ param.requires_grad = False
+
+ def forward(self, text):
+ tokens = clip.tokenize(text).to(self.device)
+ z = self.model.encode_text(tokens)
+ if self.normalize:
+ z = z / torch.linalg.norm(z, dim=1, keepdim=True)
+ return z
+
+ def encode(self, text):
+ z = self(text)
+ if z.ndim == 2:
+ z = z[:, None, :]
+ z = repeat(z, 'b 1 d -> b k d', k=self.n_repeat)
+ return z
+
+
+class FrozenClipMultiTextEmbedder(FrozenCLIPTextEmbedder):
+ def __init__(self, num_views=1, apply_all=False, **kwargs):
+ super().__init__(**kwargs)
+ self.num_views = num_views
+ self.apply_all = apply_all
+
+ def encode(self, text):
+ z = self(text)
+ if z.ndim == 2:
+ z = z[:, None, :]
+
+ if not self.apply_all:
+ new_z = torch.zeros(*z.shape[:2], z.shape[2] * self.num_views, device=z.device)
+ new_z[:, :, self.num_views // 2 * z.shape[2]: (self.num_views // 2 + 1) * z.shape[2]] = z
+ else:
+ new_z = repeat(z, 'b 1 d -> b 1 (d m)', m=self.num_views)
+
+ return new_z
+
+
+class FrozenClipImageEmbedder(nn.Module):
+ """
+ Uses the CLIP image encoder.
+ """
+
+ def __init__(
+ self,
+ model,
+ jit=False,
+ device='cuda' if torch.cuda.is_available() else 'cpu',
+ antialias=False,
+ ):
+ super().__init__()
+ self.model, _ = clip.load(name=model, device='cpu', jit=jit)
+ self.init()
+
+ self.antialias = antialias
+
+ self.register_buffer('mean', torch.Tensor([0.48145466, 0.4578275, 0.40821073]), persistent=False)
+ self.register_buffer('std', torch.Tensor([0.26862954, 0.26130258, 0.27577711]), persistent=False)
+
+ def init(self):
+ for param in self.model.parameters():
+ param.requires_grad = False
+ self.model.eval()
+
+ def preprocess(self, x):
+ x = kornia.geometry.resize(x, (224, 224),
+ interpolation='bicubic', align_corners=True,
+ antialias=self.antialias)
+ # x = (x + 1.) / 2.
+
+ # renormalize according to clip
+ x = kornia.enhance.normalize(x, self.mean, self.std)
+ return x
+
+ def forward(self, x):
+ # x is assumed to be in range [0,1]
+ return self.model.encode_image(self.preprocess(x))
+
+
+class FrozenClipMultiImageEmbedder(FrozenClipImageEmbedder):
+ """
+ Uses the CLIP image encoder with multi-image as input.
+ """
+
+ def __init__(self, num_views=1, split_per_view=1, img_dim=768, out_dim=512, key='camera', **kwargs):
+ super().__init__(**kwargs)
+ self.split_per_view = split_per_view
+ self.key = key
+ self.linear = nn.Linear(img_dim, out_dim)
+ self.view_embedding = nn.Parameter(img_dim ** -0.5 * torch.randn((1, num_views * split_per_view, img_dim)))
+
+ def forward(self, x):
+ # x is assumed to be in range [0,1]
+ if isinstance(x, torch.Tensor) and x.ndim == 5:
+ x = x.permute(1, 0, 2, 3, 4)
+ elif isinstance(x, dict):
+ x = x[self.key]
+ elif isinstance(x, torch.Tensor) and x.ndim == 3:
+ x = self.linear(x)
+ return x
+
+ with torch.no_grad():
+ img_feats = [self.model.encode_image(self.preprocess(img))[:, None] for img in x]
+ x = torch.cat(img_feats, 1).float() + self.view_embedding
+ x = self.linear(x)
+
+ return x
+
+
+class FrozenClipImagePatchEmbedder(nn.Module):
+ """
+ Uses the CLIP image encoder.
+ """
+
+ def __init__(
+ self,
+ model,
+ jit=False,
+ device='cuda' if torch.cuda.is_available() else 'cpu',
+ antialias=False,
+ img_dim=1024,
+ out_dim=512,
+ num_views=1,
+ split_per_view=1
+ ):
+ super().__init__()
+ self.model, _ = clip.load(name=model, device='cpu', jit=jit)
+ self.init()
+
+ self.antialias = antialias
+
+ self.register_buffer('mean', torch.Tensor([0.48145466, 0.4578275, 0.40821073]), persistent=False)
+ self.register_buffer('std', torch.Tensor([0.26862954, 0.26130258, 0.27577711]), persistent=False)
+ self.view_embedding = nn.Parameter(img_dim ** -0.5 * torch.randn((1, num_views * split_per_view, 1, img_dim)))
+
+ self.linear = nn.Linear(img_dim, out_dim)
+
+ def init(self):
+ for param in self.model.parameters():
+ param.requires_grad = False
+ self.model.eval()
+
+ def preprocess(self, x):
+ x = kornia.geometry.resize(x, (224, 224),
+ interpolation='bicubic', align_corners=True,
+ antialias=self.antialias)
+ # x = (x + 1.) / 2.
+
+ # renormalize according to clip
+ x = kornia.enhance.normalize(x, self.mean, self.std)
+ return x
+
+ def encode_image_patch(self, x):
+ visual_encoder = self.model.visual
+ x = x.type(self.model.dtype)
+ x = visual_encoder.conv1(x) # shape = [*, width, grid, grid]
+ x = x.reshape(x.shape[0], x.shape[1], -1) # shape = [*, width, grid ** 2]
+ x = x.permute(0, 2, 1) # shape = [*, grid ** 2, width]
+ x = torch.cat([visual_encoder.class_embedding.to(x.dtype) + torch.zeros(x.shape[0], 1, x.shape[-1], dtype=x.dtype, device=x.device), x], dim=1) # shape = [*, grid ** 2 + 1, width]
+ x = x + visual_encoder.positional_embedding.to(x.dtype)
+ x = visual_encoder.ln_pre(x)
+
+ x = x.permute(1, 0, 2) # NLD -> LND
+ x = visual_encoder.transformer(x)
+ x = x.permute(1, 0, 2) # LND -> NLD
+ x = x[:, 1:, :]
+
+ return x
+
+ def forward(self, x):
+ # x is assumed to be in range [0,1]
+ img_feats = [self.encode_image_patch(self.preprocess(img))[:, None] for img in x]
+ x = torch.cat(img_feats, 1).float() + self.view_embedding
+ x = rearrange(x, 'b v n c -> b (v n) c')
+ x = self.linear(x)
+ return x
diff --git a/lidm/modules/image_degradation/__init__.py b/lidm/modules/image_degradation/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..0db8d3e862120b2c4587c6b904c0c1153ebcdf48
--- /dev/null
+++ b/lidm/modules/image_degradation/__init__.py
@@ -0,0 +1,2 @@
+from .bsrgan import degradation_bsrgan_variant as degradation_fn_bsr
+from .bsrgan_light import degradation_bsrgan_variant as degradation_fn_bsr_light
diff --git a/lidm/modules/image_degradation/bsrgan.py b/lidm/modules/image_degradation/bsrgan.py
new file mode 100644
index 0000000000000000000000000000000000000000..5b28874c1cc2aeb47f5f5d067397272401e59dd6
--- /dev/null
+++ b/lidm/modules/image_degradation/bsrgan.py
@@ -0,0 +1,730 @@
+# -*- coding: utf-8 -*-
+"""
+# --------------------------------------------
+# Super-Resolution
+# --------------------------------------------
+#
+# Kai Zhang (cskaizhang@gmail.com)
+# https://github.com/cszn
+# From 2019/03--2021/08
+# --------------------------------------------
+"""
+
+import numpy as np
+import cv2
+import torch
+
+from functools import partial
+import random
+from scipy import ndimage
+import scipy
+import scipy.stats as ss
+from scipy.interpolate import interp2d
+from scipy.linalg import orth
+import albumentations
+
+from . import utils_image as util
+
+
+def modcrop_np(img, sf):
+ '''
+ Args:
+ img: numpy image, WxH or WxHxC
+ sf: scale factor
+ Return:
+ cropped image
+ '''
+ w, h = img.shape[:2]
+ im = np.copy(img)
+ return im[:w - w % sf, :h - h % sf, ...]
+
+
+"""
+# --------------------------------------------
+# anisotropic Gaussian kernels
+# --------------------------------------------
+"""
+
+
+def analytic_kernel(k):
+ """Calculate the X4 kernel from the X2 kernel (for proof see appendix in paper)"""
+ k_size = k.shape[0]
+ # Calculate the big kernels size
+ big_k = np.zeros((3 * k_size - 2, 3 * k_size - 2))
+ # Loop over the small kernel to fill the big one
+ for r in range(k_size):
+ for c in range(k_size):
+ big_k[2 * r:2 * r + k_size, 2 * c:2 * c + k_size] += k[r, c] * k
+ # Crop the edges of the big kernel to ignore very small values and increase run time of SR
+ crop = k_size // 2
+ cropped_big_k = big_k[crop:-crop, crop:-crop]
+ # Normalize to 1
+ return cropped_big_k / cropped_big_k.sum()
+
+
+def anisotropic_Gaussian(ksize=15, theta=np.pi, l1=6, l2=6):
+ """ generate an anisotropic Gaussian kernel
+ Args:
+ ksize : e.g., 15, kernel size
+ theta : [0, pi], rotation angle range
+ l1 : [0.1,50], scaling of eigenvalues
+ l2 : [0.1,l1], scaling of eigenvalues
+ If l1 = l2, will get an isotropic Gaussian kernel.
+ Returns:
+ k : kernel
+ """
+
+ v = np.dot(np.array([[np.cos(theta), -np.sin(theta)], [np.sin(theta), np.cos(theta)]]), np.array([1., 0.]))
+ V = np.array([[v[0], v[1]], [v[1], -v[0]]])
+ D = np.array([[l1, 0], [0, l2]])
+ Sigma = np.dot(np.dot(V, D), np.linalg.inv(V))
+ k = gm_blur_kernel(mean=[0, 0], cov=Sigma, size=ksize)
+
+ return k
+
+
+def gm_blur_kernel(mean, cov, size=15):
+ center = size / 2.0 + 0.5
+ k = np.zeros([size, size])
+ for y in range(size):
+ for x in range(size):
+ cy = y - center + 1
+ cx = x - center + 1
+ k[y, x] = ss.multivariate_normal.pdf([cx, cy], mean=mean, cov=cov)
+
+ k = k / np.sum(k)
+ return k
+
+
+def shift_pixel(x, sf, upper_left=True):
+ """shift pixel for super-resolution with different scale factors
+ Args:
+ x: WxHxC or WxH
+ sf: scale factor
+ upper_left: shift direction
+ """
+ h, w = x.shape[:2]
+ shift = (sf - 1) * 0.5
+ xv, yv = np.arange(0, w, 1.0), np.arange(0, h, 1.0)
+ if upper_left:
+ x1 = xv + shift
+ y1 = yv + shift
+ else:
+ x1 = xv - shift
+ y1 = yv - shift
+
+ x1 = np.clip(x1, 0, w - 1)
+ y1 = np.clip(y1, 0, h - 1)
+
+ if x.ndim == 2:
+ x = interp2d(xv, yv, x)(x1, y1)
+ if x.ndim == 3:
+ for i in range(x.shape[-1]):
+ x[:, :, i] = interp2d(xv, yv, x[:, :, i])(x1, y1)
+
+ return x
+
+
+def blur(x, k):
+ '''
+ x: image, NxcxHxW
+ k: kernel, Nx1xhxw
+ '''
+ n, c = x.shape[:2]
+ p1, p2 = (k.shape[-2] - 1) // 2, (k.shape[-1] - 1) // 2
+ x = torch.nn.functional.pad(x, pad=(p1, p2, p1, p2), mode='replicate')
+ k = k.repeat(1, c, 1, 1)
+ k = k.view(-1, 1, k.shape[2], k.shape[3])
+ x = x.view(1, -1, x.shape[2], x.shape[3])
+ x = torch.nn.functional.conv2d(x, k, bias=None, stride=1, padding=0, groups=n * c)
+ x = x.view(n, c, x.shape[2], x.shape[3])
+
+ return x
+
+
+def gen_kernel(k_size=np.array([15, 15]), scale_factor=np.array([4, 4]), min_var=0.6, max_var=10., noise_level=0):
+ """"
+ # modified version of https://github.com/assafshocher/BlindSR_dataset_generator
+ # Kai Zhang
+ # min_var = 0.175 * sf # variance of the gaussian kernel will be sampled between min_var and max_var
+ # max_var = 2.5 * sf
+ """
+ # Set random eigen-vals (lambdas) and angle (theta) for COV matrix
+ lambda_1 = min_var + np.random.rand() * (max_var - min_var)
+ lambda_2 = min_var + np.random.rand() * (max_var - min_var)
+ theta = np.random.rand() * np.pi # random theta
+ noise = -noise_level + np.random.rand(*k_size) * noise_level * 2
+
+ # Set COV matrix using Lambdas and Theta
+ LAMBDA = np.diag([lambda_1, lambda_2])
+ Q = np.array([[np.cos(theta), -np.sin(theta)],
+ [np.sin(theta), np.cos(theta)]])
+ SIGMA = Q @ LAMBDA @ Q.T
+ INV_SIGMA = np.linalg.inv(SIGMA)[None, None, :, :]
+
+ # Set expectation position (shifting kernel for aligned image)
+ MU = k_size // 2 - 0.5 * (scale_factor - 1) # - 0.5 * (scale_factor - k_size % 2)
+ MU = MU[None, None, :, None]
+
+ # Create meshgrid for Gaussian
+ [X, Y] = np.meshgrid(range(k_size[0]), range(k_size[1]))
+ Z = np.stack([X, Y], 2)[:, :, :, None]
+
+ # Calcualte Gaussian for every pixel of the kernel
+ ZZ = Z - MU
+ ZZ_t = ZZ.transpose(0, 1, 3, 2)
+ raw_kernel = np.exp(-0.5 * np.squeeze(ZZ_t @ INV_SIGMA @ ZZ)) * (1 + noise)
+
+ # shift the kernel so it will be centered
+ # raw_kernel_centered = kernel_shift(raw_kernel, scale_factor)
+
+ # Normalize the kernel and return
+ # kernel = raw_kernel_centered / np.sum(raw_kernel_centered)
+ kernel = raw_kernel / np.sum(raw_kernel)
+ return kernel
+
+
+def fspecial_gaussian(hsize, sigma):
+ hsize = [hsize, hsize]
+ siz = [(hsize[0] - 1.0) / 2.0, (hsize[1] - 1.0) / 2.0]
+ std = sigma
+ [x, y] = np.meshgrid(np.arange(-siz[1], siz[1] + 1), np.arange(-siz[0], siz[0] + 1))
+ arg = -(x * x + y * y) / (2 * std * std)
+ h = np.exp(arg)
+ h[h < scipy.finfo(float).eps * h.max()] = 0
+ sumh = h.sum()
+ if sumh != 0:
+ h = h / sumh
+ return h
+
+
+def fspecial_laplacian(alpha):
+ alpha = max([0, min([alpha, 1])])
+ h1 = alpha / (alpha + 1)
+ h2 = (1 - alpha) / (alpha + 1)
+ h = [[h1, h2, h1], [h2, -4 / (alpha + 1), h2], [h1, h2, h1]]
+ h = np.array(h)
+ return h
+
+
+def fspecial(filter_type, *args, **kwargs):
+ '''
+ python code from:
+ https://github.com/ronaldosena/imagens-medicas-2/blob/40171a6c259edec7827a6693a93955de2bd39e76/Aulas/aula_2_-_uniform_filter/matlab_fspecial.py
+ '''
+ if filter_type == 'gaussian':
+ return fspecial_gaussian(*args, **kwargs)
+ if filter_type == 'laplacian':
+ return fspecial_laplacian(*args, **kwargs)
+
+
+"""
+# --------------------------------------------
+# degradation models
+# --------------------------------------------
+"""
+
+
+def bicubic_degradation(x, sf=3):
+ '''
+ Args:
+ x: HxWxC image, [0, 1]
+ sf: down-scale factor
+ Return:
+ bicubicly downsampled LR image
+ '''
+ x = util.imresize_np(x, scale=1 / sf)
+ return x
+
+
+def srmd_degradation(x, k, sf=3):
+ ''' blur + bicubic downsampling
+ Args:
+ x: HxWxC image, [0, 1]
+ k: hxw, double
+ sf: down-scale factor
+ Return:
+ downsampled LR image
+ Reference:
+ @inproceedings{zhang2018learning,
+ title={Learning a single convolutional super-resolution network for multiple degradations},
+ author={Zhang, Kai and Zuo, Wangmeng and Zhang, Lei},
+ booktitle={IEEE Conference on Computer Vision and Pattern Recognition},
+ pages={3262--3271},
+ year={2018}
+ }
+ '''
+ x = ndimage.filters.convolve(x, np.expand_dims(k, axis=2), mode='wrap') # 'nearest' | 'mirror'
+ x = bicubic_degradation(x, sf=sf)
+ return x
+
+
+def dpsr_degradation(x, k, sf=3):
+ ''' bicubic downsampling + blur
+ Args:
+ x: HxWxC image, [0, 1]
+ k: hxw, double
+ sf: down-scale factor
+ Return:
+ downsampled LR image
+ Reference:
+ @inproceedings{zhang2019deep,
+ title={Deep Plug-and-Play Super-Resolution for Arbitrary Blur Kernels},
+ author={Zhang, Kai and Zuo, Wangmeng and Zhang, Lei},
+ booktitle={IEEE Conference on Computer Vision and Pattern Recognition},
+ pages={1671--1681},
+ year={2019}
+ }
+ '''
+ x = bicubic_degradation(x, sf=sf)
+ x = ndimage.filters.convolve(x, np.expand_dims(k, axis=2), mode='wrap')
+ return x
+
+
+def classical_degradation(x, k, sf=3):
+ ''' blur + downsampling
+ Args:
+ x: HxWxC image, [0, 1]/[0, 255]
+ k: hxw, double
+ sf: down-scale factor
+ Return:
+ downsampled LR image
+ '''
+ x = ndimage.filters.convolve(x, np.expand_dims(k, axis=2), mode='wrap')
+ # x = filters.correlate(x, np.expand_dims(np.flip(k), axis=2))
+ st = 0
+ return x[st::sf, st::sf, ...]
+
+
+def add_sharpening(img, weight=0.5, radius=50, threshold=10):
+ """USM sharpening. borrowed from real-ESRGAN
+ Input image: I; Blurry image: B.
+ 1. K = I + weight * (I - B)
+ 2. Mask = 1 if abs(I - B) > threshold, else: 0
+ 3. Blur mask:
+ 4. Out = Mask * K + (1 - Mask) * I
+ Args:
+ img (Numpy array): Input image, HWC, BGR; float32, [0, 1].
+ weight (float): Sharp weight. Default: 1.
+ radius (float): Kernel size of Gaussian blur. Default: 50.
+ threshold (int):
+ """
+ if radius % 2 == 0:
+ radius += 1
+ blur = cv2.GaussianBlur(img, (radius, radius), 0)
+ residual = img - blur
+ mask = np.abs(residual) * 255 > threshold
+ mask = mask.astype('float32')
+ soft_mask = cv2.GaussianBlur(mask, (radius, radius), 0)
+
+ K = img + weight * residual
+ K = np.clip(K, 0, 1)
+ return soft_mask * K + (1 - soft_mask) * img
+
+
+def add_blur(img, sf=4):
+ wd2 = 4.0 + sf
+ wd = 2.0 + 0.2 * sf
+ if random.random() < 0.5:
+ l1 = wd2 * random.random()
+ l2 = wd2 * random.random()
+ k = anisotropic_Gaussian(ksize=2 * random.randint(2, 11) + 3, theta=random.random() * np.pi, l1=l1, l2=l2)
+ else:
+ k = fspecial('gaussian', 2 * random.randint(2, 11) + 3, wd * random.random())
+ img = ndimage.filters.convolve(img, np.expand_dims(k, axis=2), mode='mirror')
+
+ return img
+
+
+def add_resize(img, sf=4):
+ rnum = np.random.rand()
+ if rnum > 0.8: # up
+ sf1 = random.uniform(1, 2)
+ elif rnum < 0.7: # down
+ sf1 = random.uniform(0.5 / sf, 1)
+ else:
+ sf1 = 1.0
+ img = cv2.resize(img, (int(sf1 * img.shape[1]), int(sf1 * img.shape[0])), interpolation=random.choice([1, 2, 3]))
+ img = np.clip(img, 0.0, 1.0)
+
+ return img
+
+
+# def add_Gaussian_noise(img, noise_level1=2, noise_level2=25):
+# noise_level = random.randint(noise_level1, noise_level2)
+# rnum = np.random.rand()
+# if rnum > 0.6: # add color Gaussian noise
+# img += np.random.normal(0, noise_level / 255.0, img.shape).astype(np.float32)
+# elif rnum < 0.4: # add grayscale Gaussian noise
+# img += np.random.normal(0, noise_level / 255.0, (*img.shape[:2], 1)).astype(np.float32)
+# else: # add noise
+# L = noise_level2 / 255.
+# D = np.diag(np.random.rand(3))
+# U = orth(np.random.rand(3, 3))
+# conv = np.dot(np.dot(np.transpose(U), D), U)
+# img += np.random.multivariate_normal([0, 0, 0], np.abs(L ** 2 * conv), img.shape[:2]).astype(np.float32)
+# img = np.clip(img, 0.0, 1.0)
+# return img
+
+def add_Gaussian_noise(img, noise_level1=2, noise_level2=25):
+ noise_level = random.randint(noise_level1, noise_level2)
+ rnum = np.random.rand()
+ if rnum > 0.6: # add color Gaussian noise
+ img = img + np.random.normal(0, noise_level / 255.0, img.shape).astype(np.float32)
+ elif rnum < 0.4: # add grayscale Gaussian noise
+ img = img + np.random.normal(0, noise_level / 255.0, (*img.shape[:2], 1)).astype(np.float32)
+ else: # add noise
+ L = noise_level2 / 255.
+ D = np.diag(np.random.rand(3))
+ U = orth(np.random.rand(3, 3))
+ conv = np.dot(np.dot(np.transpose(U), D), U)
+ img = img + np.random.multivariate_normal([0, 0, 0], np.abs(L ** 2 * conv), img.shape[:2]).astype(np.float32)
+ img = np.clip(img, 0.0, 1.0)
+ return img
+
+
+def add_speckle_noise(img, noise_level1=2, noise_level2=25):
+ noise_level = random.randint(noise_level1, noise_level2)
+ img = np.clip(img, 0.0, 1.0)
+ rnum = random.random()
+ if rnum > 0.6:
+ img += img * np.random.normal(0, noise_level / 255.0, img.shape).astype(np.float32)
+ elif rnum < 0.4:
+ img += img * np.random.normal(0, noise_level / 255.0, (*img.shape[:2], 1)).astype(np.float32)
+ else:
+ L = noise_level2 / 255.
+ D = np.diag(np.random.rand(3))
+ U = orth(np.random.rand(3, 3))
+ conv = np.dot(np.dot(np.transpose(U), D), U)
+ img += img * np.random.multivariate_normal([0, 0, 0], np.abs(L ** 2 * conv), img.shape[:2]).astype(np.float32)
+ img = np.clip(img, 0.0, 1.0)
+ return img
+
+
+def add_Poisson_noise(img):
+ img = np.clip((img * 255.0).round(), 0, 255) / 255.
+ vals = 10 ** (2 * random.random() + 2.0) # [2, 4]
+ if random.random() < 0.5:
+ img = np.random.poisson(img * vals).astype(np.float32) / vals
+ else:
+ img_gray = np.dot(img[..., :3], [0.299, 0.587, 0.114])
+ img_gray = np.clip((img_gray * 255.0).round(), 0, 255) / 255.
+ noise_gray = np.random.poisson(img_gray * vals).astype(np.float32) / vals - img_gray
+ img += noise_gray[:, :, np.newaxis]
+ img = np.clip(img, 0.0, 1.0)
+ return img
+
+
+def add_JPEG_noise(img):
+ quality_factor = random.randint(30, 95)
+ img = cv2.cvtColor(util.single2uint(img), cv2.COLOR_RGB2BGR)
+ result, encimg = cv2.imencode('.jpg', img, [int(cv2.IMWRITE_JPEG_QUALITY), quality_factor])
+ img = cv2.imdecode(encimg, 1)
+ img = cv2.cvtColor(util.uint2single(img), cv2.COLOR_BGR2RGB)
+ return img
+
+
+def random_crop(lq, hq, sf=4, lq_patchsize=64):
+ h, w = lq.shape[:2]
+ rnd_h = random.randint(0, h - lq_patchsize)
+ rnd_w = random.randint(0, w - lq_patchsize)
+ lq = lq[rnd_h:rnd_h + lq_patchsize, rnd_w:rnd_w + lq_patchsize, :]
+
+ rnd_h_H, rnd_w_H = int(rnd_h * sf), int(rnd_w * sf)
+ hq = hq[rnd_h_H:rnd_h_H + lq_patchsize * sf, rnd_w_H:rnd_w_H + lq_patchsize * sf, :]
+ return lq, hq
+
+
+def degradation_bsrgan(img, sf=4, lq_patchsize=72, isp_model=None):
+ """
+ This is the degradation model of BSRGAN from the paper
+ "Designing a Practical Degradation Model for Deep Blind Image Super-Resolution"
+ ----------
+ img: HXWXC, [0, 1], its size should be large than (lq_patchsizexsf)x(lq_patchsizexsf)
+ sf: scale factor
+ isp_model: camera ISP model
+ Returns
+ -------
+ img: low-quality patch, size: lq_patchsizeXlq_patchsizeXC, range: [0, 1]
+ hq: corresponding high-quality patch, size: (lq_patchsizexsf)X(lq_patchsizexsf)XC, range: [0, 1]
+ """
+ isp_prob, jpeg_prob, scale2_prob = 0.25, 0.9, 0.25
+ sf_ori = sf
+
+ h1, w1 = img.shape[:2]
+ img = img.copy()[:w1 - w1 % sf, :h1 - h1 % sf, ...] # mod crop
+ h, w = img.shape[:2]
+
+ if h < lq_patchsize * sf or w < lq_patchsize * sf:
+ raise ValueError(f'img size ({h1}X{w1}) is too small!')
+
+ hq = img.copy()
+
+ if sf == 4 and random.random() < scale2_prob: # downsample1
+ if np.random.rand() < 0.5:
+ img = cv2.resize(img, (int(1 / 2 * img.shape[1]), int(1 / 2 * img.shape[0])),
+ interpolation=random.choice([1, 2, 3]))
+ else:
+ img = util.imresize_np(img, 1 / 2, True)
+ img = np.clip(img, 0.0, 1.0)
+ sf = 2
+
+ shuffle_order = random.sample(range(7), 7)
+ idx1, idx2 = shuffle_order.index(2), shuffle_order.index(3)
+ if idx1 > idx2: # keep downsample3 last
+ shuffle_order[idx1], shuffle_order[idx2] = shuffle_order[idx2], shuffle_order[idx1]
+
+ for i in shuffle_order:
+
+ if i == 0:
+ img = add_blur(img, sf=sf)
+
+ elif i == 1:
+ img = add_blur(img, sf=sf)
+
+ elif i == 2:
+ a, b = img.shape[1], img.shape[0]
+ # downsample2
+ if random.random() < 0.75:
+ sf1 = random.uniform(1, 2 * sf)
+ img = cv2.resize(img, (int(1 / sf1 * img.shape[1]), int(1 / sf1 * img.shape[0])),
+ interpolation=random.choice([1, 2, 3]))
+ else:
+ k = fspecial('gaussian', 25, random.uniform(0.1, 0.6 * sf))
+ k_shifted = shift_pixel(k, sf)
+ k_shifted = k_shifted / k_shifted.sum() # blur with shifted kernel
+ img = ndimage.filters.convolve(img, np.expand_dims(k_shifted, axis=2), mode='mirror')
+ img = img[0::sf, 0::sf, ...] # nearest downsampling
+ img = np.clip(img, 0.0, 1.0)
+
+ elif i == 3:
+ # downsample3
+ img = cv2.resize(img, (int(1 / sf * a), int(1 / sf * b)), interpolation=random.choice([1, 2, 3]))
+ img = np.clip(img, 0.0, 1.0)
+
+ elif i == 4:
+ # add Gaussian noise
+ img = add_Gaussian_noise(img, noise_level1=2, noise_level2=25)
+
+ elif i == 5:
+ # add JPEG noise
+ if random.random() < jpeg_prob:
+ img = add_JPEG_noise(img)
+
+ elif i == 6:
+ # add processed camera sensor noise
+ if random.random() < isp_prob and isp_model is not None:
+ with torch.no_grad():
+ img, hq = isp_model.forward(img.copy(), hq)
+
+ # add final JPEG compression noise
+ img = add_JPEG_noise(img)
+
+ # random crop
+ img, hq = random_crop(img, hq, sf_ori, lq_patchsize)
+
+ return img, hq
+
+
+# todo no isp_model?
+def degradation_bsrgan_variant(image, sf=4, isp_model=None):
+ """
+ This is the degradation model of BSRGAN from the paper
+ "Designing a Practical Degradation Model for Deep Blind Image Super-Resolution"
+ ----------
+ sf: scale factor
+ isp_model: camera ISP model
+ Returns
+ -------
+ img: low-quality patch, size: lq_patchsizeXlq_patchsizeXC, range: [0, 1]
+ hq: corresponding high-quality patch, size: (lq_patchsizexsf)X(lq_patchsizexsf)XC, range: [0, 1]
+ """
+ image = util.uint2single(image)
+ isp_prob, jpeg_prob, scale2_prob = 0.25, 0.9, 0.25
+ sf_ori = sf
+
+ h1, w1 = image.shape[:2]
+ image = image.copy()[:w1 - w1 % sf, :h1 - h1 % sf, ...] # mod crop
+ h, w = image.shape[:2]
+
+ hq = image.copy()
+
+ if sf == 4 and random.random() < scale2_prob: # downsample1
+ if np.random.rand() < 0.5:
+ image = cv2.resize(image, (int(1 / 2 * image.shape[1]), int(1 / 2 * image.shape[0])),
+ interpolation=random.choice([1, 2, 3]))
+ else:
+ image = util.imresize_np(image, 1 / 2, True)
+ image = np.clip(image, 0.0, 1.0)
+ sf = 2
+
+ shuffle_order = random.sample(range(7), 7)
+ idx1, idx2 = shuffle_order.index(2), shuffle_order.index(3)
+ if idx1 > idx2: # keep downsample3 last
+ shuffle_order[idx1], shuffle_order[idx2] = shuffle_order[idx2], shuffle_order[idx1]
+
+ for i in shuffle_order:
+
+ if i == 0:
+ image = add_blur(image, sf=sf)
+
+ elif i == 1:
+ image = add_blur(image, sf=sf)
+
+ elif i == 2:
+ a, b = image.shape[1], image.shape[0]
+ # downsample2
+ if random.random() < 0.75:
+ sf1 = random.uniform(1, 2 * sf)
+ image = cv2.resize(image, (int(1 / sf1 * image.shape[1]), int(1 / sf1 * image.shape[0])),
+ interpolation=random.choice([1, 2, 3]))
+ else:
+ k = fspecial('gaussian', 25, random.uniform(0.1, 0.6 * sf))
+ k_shifted = shift_pixel(k, sf)
+ k_shifted = k_shifted / k_shifted.sum() # blur with shifted kernel
+ image = ndimage.filters.convolve(image, np.expand_dims(k_shifted, axis=2), mode='mirror')
+ image = image[0::sf, 0::sf, ...] # nearest downsampling
+ image = np.clip(image, 0.0, 1.0)
+
+ elif i == 3:
+ # downsample3
+ image = cv2.resize(image, (int(1 / sf * a), int(1 / sf * b)), interpolation=random.choice([1, 2, 3]))
+ image = np.clip(image, 0.0, 1.0)
+
+ elif i == 4:
+ # add Gaussian noise
+ image = add_Gaussian_noise(image, noise_level1=2, noise_level2=25)
+
+ elif i == 5:
+ # add JPEG noise
+ if random.random() < jpeg_prob:
+ image = add_JPEG_noise(image)
+
+ # elif i == 6:
+ # # add processed camera sensor noise
+ # if random.random() < isp_prob and isp_model is not None:
+ # with torch.no_grad():
+ # img, hq = isp_model.forward(img.copy(), hq)
+
+ # add final JPEG compression noise
+ image = add_JPEG_noise(image)
+ image = util.single2uint(image)
+ example = {"image":image}
+ return example
+
+
+# TODO incase there is a pickle error one needs to replace a += x with a = a + x in add_speckle_noise etc...
+def degradation_bsrgan_plus(img, sf=4, shuffle_prob=0.5, use_sharp=True, lq_patchsize=64, isp_model=None):
+ """
+ This is an extended degradation model by combining
+ the degradation models of BSRGAN and Real-ESRGAN
+ ----------
+ img: HXWXC, [0, 1], its size should be large than (lq_patchsizexsf)x(lq_patchsizexsf)
+ sf: scale factor
+ use_shuffle: the degradation shuffle
+ use_sharp: sharpening the img
+ Returns
+ -------
+ img: low-quality patch, size: lq_patchsizeXlq_patchsizeXC, range: [0, 1]
+ hq: corresponding high-quality patch, size: (lq_patchsizexsf)X(lq_patchsizexsf)XC, range: [0, 1]
+ """
+
+ h1, w1 = img.shape[:2]
+ img = img.copy()[:w1 - w1 % sf, :h1 - h1 % sf, ...] # mod crop
+ h, w = img.shape[:2]
+
+ if h < lq_patchsize * sf or w < lq_patchsize * sf:
+ raise ValueError(f'img size ({h1}X{w1}) is too small!')
+
+ if use_sharp:
+ img = add_sharpening(img)
+ hq = img.copy()
+
+ if random.random() < shuffle_prob:
+ shuffle_order = random.sample(range(13), 13)
+ else:
+ shuffle_order = list(range(13))
+ # local shuffle for noise, JPEG is always the last one
+ shuffle_order[2:6] = random.sample(shuffle_order[2:6], len(range(2, 6)))
+ shuffle_order[9:13] = random.sample(shuffle_order[9:13], len(range(9, 13)))
+
+ poisson_prob, speckle_prob, isp_prob = 0.1, 0.1, 0.1
+
+ for i in shuffle_order:
+ if i == 0:
+ img = add_blur(img, sf=sf)
+ elif i == 1:
+ img = add_resize(img, sf=sf)
+ elif i == 2:
+ img = add_Gaussian_noise(img, noise_level1=2, noise_level2=25)
+ elif i == 3:
+ if random.random() < poisson_prob:
+ img = add_Poisson_noise(img)
+ elif i == 4:
+ if random.random() < speckle_prob:
+ img = add_speckle_noise(img)
+ elif i == 5:
+ if random.random() < isp_prob and isp_model is not None:
+ with torch.no_grad():
+ img, hq = isp_model.forward(img.copy(), hq)
+ elif i == 6:
+ img = add_JPEG_noise(img)
+ elif i == 7:
+ img = add_blur(img, sf=sf)
+ elif i == 8:
+ img = add_resize(img, sf=sf)
+ elif i == 9:
+ img = add_Gaussian_noise(img, noise_level1=2, noise_level2=25)
+ elif i == 10:
+ if random.random() < poisson_prob:
+ img = add_Poisson_noise(img)
+ elif i == 11:
+ if random.random() < speckle_prob:
+ img = add_speckle_noise(img)
+ elif i == 12:
+ if random.random() < isp_prob and isp_model is not None:
+ with torch.no_grad():
+ img, hq = isp_model.forward(img.copy(), hq)
+ else:
+ print('check the shuffle!')
+
+ # resize to desired size
+ img = cv2.resize(img, (int(1 / sf * hq.shape[1]), int(1 / sf * hq.shape[0])),
+ interpolation=random.choice([1, 2, 3]))
+
+ # add final JPEG compression noise
+ img = add_JPEG_noise(img)
+
+ # random crop
+ img, hq = random_crop(img, hq, sf, lq_patchsize)
+
+ return img, hq
+
+
+if __name__ == '__main__':
+ print("hey")
+ img = util.imread_uint('utils/test.png', 3)
+ print(img)
+ img = util.uint2single(img)
+ print(img)
+ img = img[:448, :448]
+ h = img.shape[0] // 4
+ print("resizing to", h)
+ sf = 4
+ deg_fn = partial(degradation_bsrgan_variant, sf=sf)
+ for i in range(20):
+ print(i)
+ img_lq = deg_fn(img)
+ print(img_lq)
+ img_lq_bicubic = albumentations.SmallestMaxSize(max_size=h, interpolation=cv2.INTER_CUBIC)(image=img)["image"]
+ print(img_lq.shape)
+ print("bicubic", img_lq_bicubic.shape)
+ print(img_hq.shape)
+ lq_nearest = cv2.resize(util.single2uint(img_lq), (int(sf * img_lq.shape[1]), int(sf * img_lq.shape[0])),
+ interpolation=0)
+ lq_bicubic_nearest = cv2.resize(util.single2uint(img_lq_bicubic), (int(sf * img_lq.shape[1]), int(sf * img_lq.shape[0])),
+ interpolation=0)
+ img_concat = np.concatenate([lq_bicubic_nearest, lq_nearest, util.single2uint(img_hq)], axis=1)
+ util.imsave(img_concat, str(i) + '.png')
+
+
diff --git a/lidm/modules/image_degradation/bsrgan_light.py b/lidm/modules/image_degradation/bsrgan_light.py
new file mode 100644
index 0000000000000000000000000000000000000000..38f08a613ccbc6f18cd6c78f23919c222062a7cf
--- /dev/null
+++ b/lidm/modules/image_degradation/bsrgan_light.py
@@ -0,0 +1,650 @@
+# -*- coding: utf-8 -*-
+import numpy as np
+import cv2
+import torch
+
+from functools import partial
+import random
+from scipy import ndimage
+import scipy
+import scipy.stats as ss
+from scipy.interpolate import interp2d
+from scipy.linalg import orth
+import albumentations
+
+from . import utils_image as util
+
+"""
+# --------------------------------------------
+# Super-Resolution
+# --------------------------------------------
+#
+# Kai Zhang (cskaizhang@gmail.com)
+# https://github.com/cszn
+# From 2019/03--2021/08
+# --------------------------------------------
+"""
+
+
+def modcrop_np(img, sf):
+ '''
+ Args:
+ img: numpy image, WxH or WxHxC
+ sf: scale factor
+ Return:
+ cropped image
+ '''
+ w, h = img.shape[:2]
+ im = np.copy(img)
+ return im[:w - w % sf, :h - h % sf, ...]
+
+
+"""
+# --------------------------------------------
+# anisotropic Gaussian kernels
+# --------------------------------------------
+"""
+
+
+def analytic_kernel(k):
+ """Calculate the X4 kernel from the X2 kernel (for proof see appendix in paper)"""
+ k_size = k.shape[0]
+ # Calculate the big kernels size
+ big_k = np.zeros((3 * k_size - 2, 3 * k_size - 2))
+ # Loop over the small kernel to fill the big one
+ for r in range(k_size):
+ for c in range(k_size):
+ big_k[2 * r:2 * r + k_size, 2 * c:2 * c + k_size] += k[r, c] * k
+ # Crop the edges of the big kernel to ignore very small values and increase run time of SR
+ crop = k_size // 2
+ cropped_big_k = big_k[crop:-crop, crop:-crop]
+ # Normalize to 1
+ return cropped_big_k / cropped_big_k.sum()
+
+
+def anisotropic_Gaussian(ksize=15, theta=np.pi, l1=6, l2=6):
+ """ generate an anisotropic Gaussian kernel
+ Args:
+ ksize : e.g., 15, kernel size
+ theta : [0, pi], rotation angle range
+ l1 : [0.1,50], scaling of eigenvalues
+ l2 : [0.1,l1], scaling of eigenvalues
+ If l1 = l2, will get an isotropic Gaussian kernel.
+ Returns:
+ k : kernel
+ """
+
+ v = np.dot(np.array([[np.cos(theta), -np.sin(theta)], [np.sin(theta), np.cos(theta)]]), np.array([1., 0.]))
+ V = np.array([[v[0], v[1]], [v[1], -v[0]]])
+ D = np.array([[l1, 0], [0, l2]])
+ Sigma = np.dot(np.dot(V, D), np.linalg.inv(V))
+ k = gm_blur_kernel(mean=[0, 0], cov=Sigma, size=ksize)
+
+ return k
+
+
+def gm_blur_kernel(mean, cov, size=15):
+ center = size / 2.0 + 0.5
+ k = np.zeros([size, size])
+ for y in range(size):
+ for x in range(size):
+ cy = y - center + 1
+ cx = x - center + 1
+ k[y, x] = ss.multivariate_normal.pdf([cx, cy], mean=mean, cov=cov)
+
+ k = k / np.sum(k)
+ return k
+
+
+def shift_pixel(x, sf, upper_left=True):
+ """shift pixel for super-resolution with different scale factors
+ Args:
+ x: WxHxC or WxH
+ sf: scale factor
+ upper_left: shift direction
+ """
+ h, w = x.shape[:2]
+ shift = (sf - 1) * 0.5
+ xv, yv = np.arange(0, w, 1.0), np.arange(0, h, 1.0)
+ if upper_left:
+ x1 = xv + shift
+ y1 = yv + shift
+ else:
+ x1 = xv - shift
+ y1 = yv - shift
+
+ x1 = np.clip(x1, 0, w - 1)
+ y1 = np.clip(y1, 0, h - 1)
+
+ if x.ndim == 2:
+ x = interp2d(xv, yv, x)(x1, y1)
+ if x.ndim == 3:
+ for i in range(x.shape[-1]):
+ x[:, :, i] = interp2d(xv, yv, x[:, :, i])(x1, y1)
+
+ return x
+
+
+def blur(x, k):
+ '''
+ x: image, NxcxHxW
+ k: kernel, Nx1xhxw
+ '''
+ n, c = x.shape[:2]
+ p1, p2 = (k.shape[-2] - 1) // 2, (k.shape[-1] - 1) // 2
+ x = torch.nn.functional.pad(x, pad=(p1, p2, p1, p2), mode='replicate')
+ k = k.repeat(1, c, 1, 1)
+ k = k.view(-1, 1, k.shape[2], k.shape[3])
+ x = x.view(1, -1, x.shape[2], x.shape[3])
+ x = torch.nn.functional.conv2d(x, k, bias=None, stride=1, padding=0, groups=n * c)
+ x = x.view(n, c, x.shape[2], x.shape[3])
+
+ return x
+
+
+def gen_kernel(k_size=np.array([15, 15]), scale_factor=np.array([4, 4]), min_var=0.6, max_var=10., noise_level=0):
+ """"
+ # modified version of https://github.com/assafshocher/BlindSR_dataset_generator
+ # Kai Zhang
+ # min_var = 0.175 * sf # variance of the gaussian kernel will be sampled between min_var and max_var
+ # max_var = 2.5 * sf
+ """
+ # Set random eigen-vals (lambdas) and angle (theta) for COV matrix
+ lambda_1 = min_var + np.random.rand() * (max_var - min_var)
+ lambda_2 = min_var + np.random.rand() * (max_var - min_var)
+ theta = np.random.rand() * np.pi # random theta
+ noise = -noise_level + np.random.rand(*k_size) * noise_level * 2
+
+ # Set COV matrix using Lambdas and Theta
+ LAMBDA = np.diag([lambda_1, lambda_2])
+ Q = np.array([[np.cos(theta), -np.sin(theta)],
+ [np.sin(theta), np.cos(theta)]])
+ SIGMA = Q @ LAMBDA @ Q.T
+ INV_SIGMA = np.linalg.inv(SIGMA)[None, None, :, :]
+
+ # Set expectation position (shifting kernel for aligned image)
+ MU = k_size // 2 - 0.5 * (scale_factor - 1) # - 0.5 * (scale_factor - k_size % 2)
+ MU = MU[None, None, :, None]
+
+ # Create meshgrid for Gaussian
+ [X, Y] = np.meshgrid(range(k_size[0]), range(k_size[1]))
+ Z = np.stack([X, Y], 2)[:, :, :, None]
+
+ # Calcualte Gaussian for every pixel of the kernel
+ ZZ = Z - MU
+ ZZ_t = ZZ.transpose(0, 1, 3, 2)
+ raw_kernel = np.exp(-0.5 * np.squeeze(ZZ_t @ INV_SIGMA @ ZZ)) * (1 + noise)
+
+ # shift the kernel so it will be centered
+ # raw_kernel_centered = kernel_shift(raw_kernel, scale_factor)
+
+ # Normalize the kernel and return
+ # kernel = raw_kernel_centered / np.sum(raw_kernel_centered)
+ kernel = raw_kernel / np.sum(raw_kernel)
+ return kernel
+
+
+def fspecial_gaussian(hsize, sigma):
+ hsize = [hsize, hsize]
+ siz = [(hsize[0] - 1.0) / 2.0, (hsize[1] - 1.0) / 2.0]
+ std = sigma
+ [x, y] = np.meshgrid(np.arange(-siz[1], siz[1] + 1), np.arange(-siz[0], siz[0] + 1))
+ arg = -(x * x + y * y) / (2 * std * std)
+ h = np.exp(arg)
+ h[h < scipy.finfo(float).eps * h.max()] = 0
+ sumh = h.sum()
+ if sumh != 0:
+ h = h / sumh
+ return h
+
+
+def fspecial_laplacian(alpha):
+ alpha = max([0, min([alpha, 1])])
+ h1 = alpha / (alpha + 1)
+ h2 = (1 - alpha) / (alpha + 1)
+ h = [[h1, h2, h1], [h2, -4 / (alpha + 1), h2], [h1, h2, h1]]
+ h = np.array(h)
+ return h
+
+
+def fspecial(filter_type, *args, **kwargs):
+ '''
+ python code from:
+ https://github.com/ronaldosena/imagens-medicas-2/blob/40171a6c259edec7827a6693a93955de2bd39e76/Aulas/aula_2_-_uniform_filter/matlab_fspecial.py
+ '''
+ if filter_type == 'gaussian':
+ return fspecial_gaussian(*args, **kwargs)
+ if filter_type == 'laplacian':
+ return fspecial_laplacian(*args, **kwargs)
+
+
+"""
+# --------------------------------------------
+# degradation models
+# --------------------------------------------
+"""
+
+
+def bicubic_degradation(x, sf=3):
+ '''
+ Args:
+ x: HxWxC image, [0, 1]
+ sf: down-scale factor
+ Return:
+ bicubicly downsampled LR image
+ '''
+ x = util.imresize_np(x, scale=1 / sf)
+ return x
+
+
+def srmd_degradation(x, k, sf=3):
+ ''' blur + bicubic downsampling
+ Args:
+ x: HxWxC image, [0, 1]
+ k: hxw, double
+ sf: down-scale factor
+ Return:
+ downsampled LR image
+ Reference:
+ @inproceedings{zhang2018learning,
+ title={Learning a single convolutional super-resolution network for multiple degradations},
+ author={Zhang, Kai and Zuo, Wangmeng and Zhang, Lei},
+ booktitle={IEEE Conference on Computer Vision and Pattern Recognition},
+ pages={3262--3271},
+ year={2018}
+ }
+ '''
+ x = ndimage.filters.convolve(x, np.expand_dims(k, axis=2), mode='wrap') # 'nearest' | 'mirror'
+ x = bicubic_degradation(x, sf=sf)
+ return x
+
+
+def dpsr_degradation(x, k, sf=3):
+ ''' bicubic downsampling + blur
+ Args:
+ x: HxWxC image, [0, 1]
+ k: hxw, double
+ sf: down-scale factor
+ Return:
+ downsampled LR image
+ Reference:
+ @inproceedings{zhang2019deep,
+ title={Deep Plug-and-Play Super-Resolution for Arbitrary Blur Kernels},
+ author={Zhang, Kai and Zuo, Wangmeng and Zhang, Lei},
+ booktitle={IEEE Conference on Computer Vision and Pattern Recognition},
+ pages={1671--1681},
+ year={2019}
+ }
+ '''
+ x = bicubic_degradation(x, sf=sf)
+ x = ndimage.filters.convolve(x, np.expand_dims(k, axis=2), mode='wrap')
+ return x
+
+
+def classical_degradation(x, k, sf=3):
+ ''' blur + downsampling
+ Args:
+ x: HxWxC image, [0, 1]/[0, 255]
+ k: hxw, double
+ sf: down-scale factor
+ Return:
+ downsampled LR image
+ '''
+ x = ndimage.filters.convolve(x, np.expand_dims(k, axis=2), mode='wrap')
+ # x = filters.correlate(x, np.expand_dims(np.flip(k), axis=2))
+ st = 0
+ return x[st::sf, st::sf, ...]
+
+
+def add_sharpening(img, weight=0.5, radius=50, threshold=10):
+ """USM sharpening. borrowed from real-ESRGAN
+ Input image: I; Blurry image: B.
+ 1. K = I + weight * (I - B)
+ 2. Mask = 1 if abs(I - B) > threshold, else: 0
+ 3. Blur mask:
+ 4. Out = Mask * K + (1 - Mask) * I
+ Args:
+ img (Numpy array): Input image, HWC, BGR; float32, [0, 1].
+ weight (float): Sharp weight. Default: 1.
+ radius (float): Kernel size of Gaussian blur. Default: 50.
+ threshold (int):
+ """
+ if radius % 2 == 0:
+ radius += 1
+ blur = cv2.GaussianBlur(img, (radius, radius), 0)
+ residual = img - blur
+ mask = np.abs(residual) * 255 > threshold
+ mask = mask.astype('float32')
+ soft_mask = cv2.GaussianBlur(mask, (radius, radius), 0)
+
+ K = img + weight * residual
+ K = np.clip(K, 0, 1)
+ return soft_mask * K + (1 - soft_mask) * img
+
+
+def add_blur(img, sf=4):
+ wd2 = 4.0 + sf
+ wd = 2.0 + 0.2 * sf
+
+ wd2 = wd2/4
+ wd = wd/4
+
+ if random.random() < 0.5:
+ l1 = wd2 * random.random()
+ l2 = wd2 * random.random()
+ k = anisotropic_Gaussian(ksize=random.randint(2, 11) + 3, theta=random.random() * np.pi, l1=l1, l2=l2)
+ else:
+ k = fspecial('gaussian', random.randint(2, 4) + 3, wd * random.random())
+ img = ndimage.filters.convolve(img, np.expand_dims(k, axis=2), mode='mirror')
+
+ return img
+
+
+def add_resize(img, sf=4):
+ rnum = np.random.rand()
+ if rnum > 0.8: # up
+ sf1 = random.uniform(1, 2)
+ elif rnum < 0.7: # down
+ sf1 = random.uniform(0.5 / sf, 1)
+ else:
+ sf1 = 1.0
+ img = cv2.resize(img, (int(sf1 * img.shape[1]), int(sf1 * img.shape[0])), interpolation=random.choice([1, 2, 3]))
+ img = np.clip(img, 0.0, 1.0)
+
+ return img
+
+
+# def add_Gaussian_noise(img, noise_level1=2, noise_level2=25):
+# noise_level = random.randint(noise_level1, noise_level2)
+# rnum = np.random.rand()
+# if rnum > 0.6: # add color Gaussian noise
+# img += np.random.normal(0, noise_level / 255.0, img.shape).astype(np.float32)
+# elif rnum < 0.4: # add grayscale Gaussian noise
+# img += np.random.normal(0, noise_level / 255.0, (*img.shape[:2], 1)).astype(np.float32)
+# else: # add noise
+# L = noise_level2 / 255.
+# D = np.diag(np.random.rand(3))
+# U = orth(np.random.rand(3, 3))
+# conv = np.dot(np.dot(np.transpose(U), D), U)
+# img += np.random.multivariate_normal([0, 0, 0], np.abs(L ** 2 * conv), img.shape[:2]).astype(np.float32)
+# img = np.clip(img, 0.0, 1.0)
+# return img
+
+def add_Gaussian_noise(img, noise_level1=2, noise_level2=25):
+ noise_level = random.randint(noise_level1, noise_level2)
+ rnum = np.random.rand()
+ if rnum > 0.6: # add color Gaussian noise
+ img = img + np.random.normal(0, noise_level / 255.0, img.shape).astype(np.float32)
+ elif rnum < 0.4: # add grayscale Gaussian noise
+ img = img + np.random.normal(0, noise_level / 255.0, (*img.shape[:2], 1)).astype(np.float32)
+ else: # add noise
+ L = noise_level2 / 255.
+ D = np.diag(np.random.rand(3))
+ U = orth(np.random.rand(3, 3))
+ conv = np.dot(np.dot(np.transpose(U), D), U)
+ img = img + np.random.multivariate_normal([0, 0, 0], np.abs(L ** 2 * conv), img.shape[:2]).astype(np.float32)
+ img = np.clip(img, 0.0, 1.0)
+ return img
+
+
+def add_speckle_noise(img, noise_level1=2, noise_level2=25):
+ noise_level = random.randint(noise_level1, noise_level2)
+ img = np.clip(img, 0.0, 1.0)
+ rnum = random.random()
+ if rnum > 0.6:
+ img += img * np.random.normal(0, noise_level / 255.0, img.shape).astype(np.float32)
+ elif rnum < 0.4:
+ img += img * np.random.normal(0, noise_level / 255.0, (*img.shape[:2], 1)).astype(np.float32)
+ else:
+ L = noise_level2 / 255.
+ D = np.diag(np.random.rand(3))
+ U = orth(np.random.rand(3, 3))
+ conv = np.dot(np.dot(np.transpose(U), D), U)
+ img += img * np.random.multivariate_normal([0, 0, 0], np.abs(L ** 2 * conv), img.shape[:2]).astype(np.float32)
+ img = np.clip(img, 0.0, 1.0)
+ return img
+
+
+def add_Poisson_noise(img):
+ img = np.clip((img * 255.0).round(), 0, 255) / 255.
+ vals = 10 ** (2 * random.random() + 2.0) # [2, 4]
+ if random.random() < 0.5:
+ img = np.random.poisson(img * vals).astype(np.float32) / vals
+ else:
+ img_gray = np.dot(img[..., :3], [0.299, 0.587, 0.114])
+ img_gray = np.clip((img_gray * 255.0).round(), 0, 255) / 255.
+ noise_gray = np.random.poisson(img_gray * vals).astype(np.float32) / vals - img_gray
+ img += noise_gray[:, :, np.newaxis]
+ img = np.clip(img, 0.0, 1.0)
+ return img
+
+
+def add_JPEG_noise(img):
+ quality_factor = random.randint(80, 95)
+ img = cv2.cvtColor(util.single2uint(img), cv2.COLOR_RGB2BGR)
+ result, encimg = cv2.imencode('.jpg', img, [int(cv2.IMWRITE_JPEG_QUALITY), quality_factor])
+ img = cv2.imdecode(encimg, 1)
+ img = cv2.cvtColor(util.uint2single(img), cv2.COLOR_BGR2RGB)
+ return img
+
+
+def random_crop(lq, hq, sf=4, lq_patchsize=64):
+ h, w = lq.shape[:2]
+ rnd_h = random.randint(0, h - lq_patchsize)
+ rnd_w = random.randint(0, w - lq_patchsize)
+ lq = lq[rnd_h:rnd_h + lq_patchsize, rnd_w:rnd_w + lq_patchsize, :]
+
+ rnd_h_H, rnd_w_H = int(rnd_h * sf), int(rnd_w * sf)
+ hq = hq[rnd_h_H:rnd_h_H + lq_patchsize * sf, rnd_w_H:rnd_w_H + lq_patchsize * sf, :]
+ return lq, hq
+
+
+def degradation_bsrgan(img, sf=4, lq_patchsize=72, isp_model=None):
+ """
+ This is the degradation model of BSRGAN from the paper
+ "Designing a Practical Degradation Model for Deep Blind Image Super-Resolution"
+ ----------
+ img: HXWXC, [0, 1], its size should be large than (lq_patchsizexsf)x(lq_patchsizexsf)
+ sf: scale factor
+ isp_model: camera ISP model
+ Returns
+ -------
+ img: low-quality patch, size: lq_patchsizeXlq_patchsizeXC, range: [0, 1]
+ hq: corresponding high-quality patch, size: (lq_patchsizexsf)X(lq_patchsizexsf)XC, range: [0, 1]
+ """
+ isp_prob, jpeg_prob, scale2_prob = 0.25, 0.9, 0.25
+ sf_ori = sf
+
+ h1, w1 = img.shape[:2]
+ img = img.copy()[:w1 - w1 % sf, :h1 - h1 % sf, ...] # mod crop
+ h, w = img.shape[:2]
+
+ if h < lq_patchsize * sf or w < lq_patchsize * sf:
+ raise ValueError(f'img size ({h1}X{w1}) is too small!')
+
+ hq = img.copy()
+
+ if sf == 4 and random.random() < scale2_prob: # downsample1
+ if np.random.rand() < 0.5:
+ img = cv2.resize(img, (int(1 / 2 * img.shape[1]), int(1 / 2 * img.shape[0])),
+ interpolation=random.choice([1, 2, 3]))
+ else:
+ img = util.imresize_np(img, 1 / 2, True)
+ img = np.clip(img, 0.0, 1.0)
+ sf = 2
+
+ shuffle_order = random.sample(range(7), 7)
+ idx1, idx2 = shuffle_order.index(2), shuffle_order.index(3)
+ if idx1 > idx2: # keep downsample3 last
+ shuffle_order[idx1], shuffle_order[idx2] = shuffle_order[idx2], shuffle_order[idx1]
+
+ for i in shuffle_order:
+
+ if i == 0:
+ img = add_blur(img, sf=sf)
+
+ elif i == 1:
+ img = add_blur(img, sf=sf)
+
+ elif i == 2:
+ a, b = img.shape[1], img.shape[0]
+ # downsample2
+ if random.random() < 0.75:
+ sf1 = random.uniform(1, 2 * sf)
+ img = cv2.resize(img, (int(1 / sf1 * img.shape[1]), int(1 / sf1 * img.shape[0])),
+ interpolation=random.choice([1, 2, 3]))
+ else:
+ k = fspecial('gaussian', 25, random.uniform(0.1, 0.6 * sf))
+ k_shifted = shift_pixel(k, sf)
+ k_shifted = k_shifted / k_shifted.sum() # blur with shifted kernel
+ img = ndimage.filters.convolve(img, np.expand_dims(k_shifted, axis=2), mode='mirror')
+ img = img[0::sf, 0::sf, ...] # nearest downsampling
+ img = np.clip(img, 0.0, 1.0)
+
+ elif i == 3:
+ # downsample3
+ img = cv2.resize(img, (int(1 / sf * a), int(1 / sf * b)), interpolation=random.choice([1, 2, 3]))
+ img = np.clip(img, 0.0, 1.0)
+
+ elif i == 4:
+ # add Gaussian noise
+ img = add_Gaussian_noise(img, noise_level1=2, noise_level2=8)
+
+ elif i == 5:
+ # add JPEG noise
+ if random.random() < jpeg_prob:
+ img = add_JPEG_noise(img)
+
+ elif i == 6:
+ # add processed camera sensor noise
+ if random.random() < isp_prob and isp_model is not None:
+ with torch.no_grad():
+ img, hq = isp_model.forward(img.copy(), hq)
+
+ # add final JPEG compression noise
+ img = add_JPEG_noise(img)
+
+ # random crop
+ img, hq = random_crop(img, hq, sf_ori, lq_patchsize)
+
+ return img, hq
+
+
+# todo no isp_model?
+def degradation_bsrgan_variant(image, sf=4, isp_model=None):
+ """
+ This is the degradation model of BSRGAN from the paper
+ "Designing a Practical Degradation Model for Deep Blind Image Super-Resolution"
+ ----------
+ sf: scale factor
+ isp_model: camera ISP model
+ Returns
+ -------
+ img: low-quality patch, size: lq_patchsizeXlq_patchsizeXC, range: [0, 1]
+ hq: corresponding high-quality patch, size: (lq_patchsizexsf)X(lq_patchsizexsf)XC, range: [0, 1]
+ """
+ image = util.uint2single(image)
+ isp_prob, jpeg_prob, scale2_prob = 0.25, 0.9, 0.25
+ sf_ori = sf
+
+ h1, w1 = image.shape[:2]
+ image = image.copy()[:w1 - w1 % sf, :h1 - h1 % sf, ...] # mod crop
+ h, w = image.shape[:2]
+
+ hq = image.copy()
+
+ if sf == 4 and random.random() < scale2_prob: # downsample1
+ if np.random.rand() < 0.5:
+ image = cv2.resize(image, (int(1 / 2 * image.shape[1]), int(1 / 2 * image.shape[0])),
+ interpolation=random.choice([1, 2, 3]))
+ else:
+ image = util.imresize_np(image, 1 / 2, True)
+ image = np.clip(image, 0.0, 1.0)
+ sf = 2
+
+ shuffle_order = random.sample(range(7), 7)
+ idx1, idx2 = shuffle_order.index(2), shuffle_order.index(3)
+ if idx1 > idx2: # keep downsample3 last
+ shuffle_order[idx1], shuffle_order[idx2] = shuffle_order[idx2], shuffle_order[idx1]
+
+ for i in shuffle_order:
+
+ if i == 0:
+ image = add_blur(image, sf=sf)
+
+ # elif i == 1:
+ # image = add_blur(image, sf=sf)
+
+ if i == 0:
+ pass
+
+ elif i == 2:
+ a, b = image.shape[1], image.shape[0]
+ # downsample2
+ if random.random() < 0.8:
+ sf1 = random.uniform(1, 2 * sf)
+ image = cv2.resize(image, (int(1 / sf1 * image.shape[1]), int(1 / sf1 * image.shape[0])),
+ interpolation=random.choice([1, 2, 3]))
+ else:
+ k = fspecial('gaussian', 25, random.uniform(0.1, 0.6 * sf))
+ k_shifted = shift_pixel(k, sf)
+ k_shifted = k_shifted / k_shifted.sum() # blur with shifted kernel
+ image = ndimage.filters.convolve(image, np.expand_dims(k_shifted, axis=2), mode='mirror')
+ image = image[0::sf, 0::sf, ...] # nearest downsampling
+
+ image = np.clip(image, 0.0, 1.0)
+
+ elif i == 3:
+ # downsample3
+ image = cv2.resize(image, (int(1 / sf * a), int(1 / sf * b)), interpolation=random.choice([1, 2, 3]))
+ image = np.clip(image, 0.0, 1.0)
+
+ elif i == 4:
+ # add Gaussian noise
+ image = add_Gaussian_noise(image, noise_level1=1, noise_level2=2)
+
+ elif i == 5:
+ # add JPEG noise
+ if random.random() < jpeg_prob:
+ image = add_JPEG_noise(image)
+ #
+ # elif i == 6:
+ # # add processed camera sensor noise
+ # if random.random() < isp_prob and isp_model is not None:
+ # with torch.no_grad():
+ # img, hq = isp_model.forward(img.copy(), hq)
+
+ # add final JPEG compression noise
+ image = add_JPEG_noise(image)
+ image = util.single2uint(image)
+ example = {"image": image}
+ return example
+
+
+
+
+if __name__ == '__main__':
+ print("hey")
+ img = util.imread_uint('utils/test.png', 3)
+ img = img[:448, :448]
+ h = img.shape[0] // 4
+ print("resizing to", h)
+ sf = 4
+ deg_fn = partial(degradation_bsrgan_variant, sf=sf)
+ for i in range(20):
+ print(i)
+ img_hq = img
+ img_lq = deg_fn(img)["image"]
+ img_hq, img_lq = util.uint2single(img_hq), util.uint2single(img_lq)
+ print(img_lq)
+ img_lq_bicubic = albumentations.SmallestMaxSize(max_size=h, interpolation=cv2.INTER_CUBIC)(image=img_hq)["image"]
+ print(img_lq.shape)
+ print("bicubic", img_lq_bicubic.shape)
+ print(img_hq.shape)
+ lq_nearest = cv2.resize(util.single2uint(img_lq), (int(sf * img_lq.shape[1]), int(sf * img_lq.shape[0])),
+ interpolation=0)
+ lq_bicubic_nearest = cv2.resize(util.single2uint(img_lq_bicubic),
+ (int(sf * img_lq.shape[1]), int(sf * img_lq.shape[0])),
+ interpolation=0)
+ img_concat = np.concatenate([lq_bicubic_nearest, lq_nearest, util.single2uint(img_hq)], axis=1)
+ util.imsave(img_concat, str(i) + '.png')
diff --git a/lidm/modules/image_degradation/utils/test.png b/lidm/modules/image_degradation/utils/test.png
new file mode 100644
index 0000000000000000000000000000000000000000..4249b43de0f22707758d13c240268a401642f6e6
Binary files /dev/null and b/lidm/modules/image_degradation/utils/test.png differ
diff --git a/lidm/modules/image_degradation/utils_image.py b/lidm/modules/image_degradation/utils_image.py
new file mode 100644
index 0000000000000000000000000000000000000000..0175f155ad900ae33c3c46ed87f49b352e3faf98
--- /dev/null
+++ b/lidm/modules/image_degradation/utils_image.py
@@ -0,0 +1,916 @@
+import os
+import math
+import random
+import numpy as np
+import torch
+import cv2
+from torchvision.utils import make_grid
+from datetime import datetime
+#import matplotlib.pyplot as plt # TODO: check with Dominik, also bsrgan.py vs bsrgan_light.py
+
+
+os.environ["KMP_DUPLICATE_LIB_OK"]="TRUE"
+
+
+'''
+# --------------------------------------------
+# Kai Zhang (github: https://github.com/cszn)
+# 03/Mar/2019
+# --------------------------------------------
+# https://github.com/twhui/SRGAN-pyTorch
+# https://github.com/xinntao/BasicSR
+# --------------------------------------------
+'''
+
+
+IMG_EXTENSIONS = ['.jpg', '.JPG', '.jpeg', '.JPEG', '.png', '.PNG', '.ppm', '.PPM', '.bmp', '.BMP', '.tif']
+
+
+def is_image_file(filename):
+ return any(filename.endswith(extension) for extension in IMG_EXTENSIONS)
+
+
+def get_timestamp():
+ return datetime.now().strftime('%y%m%d-%H%M%S')
+
+
+def imshow(x, title=None, cbar=False, figsize=None):
+ plt.figure(figsize=figsize)
+ plt.imshow(np.squeeze(x), interpolation='nearest', cmap='gray')
+ if title:
+ plt.title(title)
+ if cbar:
+ plt.colorbar()
+ plt.show()
+
+
+def surf(Z, cmap='rainbow', figsize=None):
+ plt.figure(figsize=figsize)
+ ax3 = plt.axes(projection='3d')
+
+ w, h = Z.shape[:2]
+ xx = np.arange(0,w,1)
+ yy = np.arange(0,h,1)
+ X, Y = np.meshgrid(xx, yy)
+ ax3.plot_surface(X,Y,Z,cmap=cmap)
+ #ax3.contour(X,Y,Z, zdim='z',offset=-2，cmap=cmap)
+ plt.show()
+
+
+'''
+# --------------------------------------------
+# get image pathes
+# --------------------------------------------
+'''
+
+
+def get_image_paths(dataroot):
+ paths = None # return None if dataroot is None
+ if dataroot is not None:
+ paths = sorted(_get_paths_from_images(dataroot))
+ return paths
+
+
+def _get_paths_from_images(path):
+ assert os.path.isdir(path), '{:s} is not a valid directory'.format(path)
+ images = []
+ for dirpath, _, fnames in sorted(os.walk(path)):
+ for fname in sorted(fnames):
+ if is_image_file(fname):
+ img_path = os.path.join(dirpath, fname)
+ images.append(img_path)
+ assert images, '{:s} has no valid image file'.format(path)
+ return images
+
+
+'''
+# --------------------------------------------
+# split large images into small images
+# --------------------------------------------
+'''
+
+
+def patches_from_image(img, p_size=512, p_overlap=64, p_max=800):
+ w, h = img.shape[:2]
+ patches = []
+ if w > p_max and h > p_max:
+ w1 = list(np.arange(0, w-p_size, p_size-p_overlap, dtype=np.int))
+ h1 = list(np.arange(0, h-p_size, p_size-p_overlap, dtype=np.int))
+ w1.append(w-p_size)
+ h1.append(h-p_size)
+# print(w1)
+# print(h1)
+ for i in w1:
+ for j in h1:
+ patches.append(img[i:i+p_size, j:j+p_size,:])
+ else:
+ patches.append(img)
+
+ return patches
+
+
+def imssave(imgs, img_path):
+ """
+ imgs: list, N images of size WxHxC
+ """
+ img_name, ext = os.path.splitext(os.path.basename(img_path))
+
+ for i, img in enumerate(imgs):
+ if img.ndim == 3:
+ img = img[:, :, [2, 1, 0]]
+ new_path = os.path.join(os.path.dirname(img_path), img_name+str('_s{:04d}'.format(i))+'.png')
+ cv2.imwrite(new_path, img)
+
+
+def split_imageset(original_dataroot, taget_dataroot, n_channels=3, p_size=800, p_overlap=96, p_max=1000):
+ """
+ split the large images from original_dataroot into small overlapped images with size (p_size)x(p_size),
+ and save them into taget_dataroot; only the images with larger size than (p_max)x(p_max)
+ will be splitted.
+ Args:
+ original_dataroot:
+ taget_dataroot:
+ p_size: size of small images
+ p_overlap: patch size in training is a good choice
+ p_max: images with smaller size than (p_max)x(p_max) keep unchanged.
+ """
+ paths = get_image_paths(original_dataroot)
+ for img_path in paths:
+ # img_name, ext = os.path.splitext(os.path.basename(img_path))
+ img = imread_uint(img_path, n_channels=n_channels)
+ patches = patches_from_image(img, p_size, p_overlap, p_max)
+ imssave(patches, os.path.join(taget_dataroot,os.path.basename(img_path)))
+ #if original_dataroot == taget_dataroot:
+ #del img_path
+
+'''
+# --------------------------------------------
+# makedir
+# --------------------------------------------
+'''
+
+
+def mkdir(path):
+ if not os.path.exists(path):
+ os.makedirs(path)
+
+
+def mkdirs(paths):
+ if isinstance(paths, str):
+ mkdir(paths)
+ else:
+ for path in paths:
+ mkdir(path)
+
+
+def mkdir_and_rename(path):
+ if os.path.exists(path):
+ new_name = path + '_archived_' + get_timestamp()
+ print('Path already exists. Rename it to [{:s}]'.format(new_name))
+ os.rename(path, new_name)
+ os.makedirs(path)
+
+
+'''
+# --------------------------------------------
+# read image from path
+# opencv is fast, but read BGR numpy image
+# --------------------------------------------
+'''
+
+
+# --------------------------------------------
+# get uint8 image of size HxWxn_channles (RGB)
+# --------------------------------------------
+def imread_uint(path, n_channels=3):
+ # input: path
+ # output: HxWx3(RGB or GGG), or HxWx1 (G)
+ if n_channels == 1:
+ img = cv2.imread(path, 0) # cv2.IMREAD_GRAYSCALE
+ img = np.expand_dims(img, axis=2) # HxWx1
+ elif n_channels == 3:
+ img = cv2.imread(path, cv2.IMREAD_UNCHANGED) # BGR or G
+ if img.ndim == 2:
+ img = cv2.cvtColor(img, cv2.COLOR_GRAY2RGB) # GGG
+ else:
+ img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) # RGB
+ return img
+
+
+# --------------------------------------------
+# matlab's imwrite
+# --------------------------------------------
+def imsave(img, img_path):
+ img = np.squeeze(img)
+ if img.ndim == 3:
+ img = img[:, :, [2, 1, 0]]
+ cv2.imwrite(img_path, img)
+
+def imwrite(img, img_path):
+ img = np.squeeze(img)
+ if img.ndim == 3:
+ img = img[:, :, [2, 1, 0]]
+ cv2.imwrite(img_path, img)
+
+
+
+# --------------------------------------------
+# get single image of size HxWxn_channles (BGR)
+# --------------------------------------------
+def read_img(path):
+ # read image by cv2
+ # return: Numpy float32, HWC, BGR, [0,1]
+ img = cv2.imread(path, cv2.IMREAD_UNCHANGED) # cv2.IMREAD_GRAYSCALE
+ img = img.astype(np.float32) / 255.
+ if img.ndim == 2:
+ img = np.expand_dims(img, axis=2)
+ # some images have 4 channels
+ if img.shape[2] > 3:
+ img = img[:, :, :3]
+ return img
+
+
+'''
+# --------------------------------------------
+# image format conversion
+# --------------------------------------------
+# numpy(single) <---> numpy(unit)
+# numpy(single) <---> tensor
+# numpy(unit) <---> tensor
+# --------------------------------------------
+'''
+
+
+# --------------------------------------------
+# numpy(single) [0, 1] <---> numpy(unit)
+# --------------------------------------------
+
+
+def uint2single(img):
+
+ return np.float32(img/255.)
+
+
+def single2uint(img):
+
+ return np.uint8((img.clip(0, 1)*255.).round())
+
+
+def uint162single(img):
+
+ return np.float32(img/65535.)
+
+
+def single2uint16(img):
+
+ return np.uint16((img.clip(0, 1)*65535.).round())
+
+
+# --------------------------------------------
+# numpy(unit) (HxWxC or HxW) <---> tensor
+# --------------------------------------------
+
+
+# convert uint to 4-dimensional torch tensor
+def uint2tensor4(img):
+ if img.ndim == 2:
+ img = np.expand_dims(img, axis=2)
+ return torch.from_numpy(np.ascontiguousarray(img)).permute(2, 0, 1).float().div(255.).unsqueeze(0)
+
+
+# convert uint to 3-dimensional torch tensor
+def uint2tensor3(img):
+ if img.ndim == 2:
+ img = np.expand_dims(img, axis=2)
+ return torch.from_numpy(np.ascontiguousarray(img)).permute(2, 0, 1).float().div(255.)
+
+
+# convert 2/3/4-dimensional torch tensor to uint
+def tensor2uint(img):
+ img = img.data.squeeze().float().clamp_(0, 1).cpu().numpy()
+ if img.ndim == 3:
+ img = np.transpose(img, (1, 2, 0))
+ return np.uint8((img*255.0).round())
+
+
+# --------------------------------------------
+# numpy(single) (HxWxC) <---> tensor
+# --------------------------------------------
+
+
+# convert single (HxWxC) to 3-dimensional torch tensor
+def single2tensor3(img):
+ return torch.from_numpy(np.ascontiguousarray(img)).permute(2, 0, 1).float()
+
+
+# convert single (HxWxC) to 4-dimensional torch tensor
+def single2tensor4(img):
+ return torch.from_numpy(np.ascontiguousarray(img)).permute(2, 0, 1).float().unsqueeze(0)
+
+
+# convert torch tensor to single
+def tensor2single(img):
+ img = img.data.squeeze().float().cpu().numpy()
+ if img.ndim == 3:
+ img = np.transpose(img, (1, 2, 0))
+
+ return img
+
+# convert torch tensor to single
+def tensor2single3(img):
+ img = img.data.squeeze().float().cpu().numpy()
+ if img.ndim == 3:
+ img = np.transpose(img, (1, 2, 0))
+ elif img.ndim == 2:
+ img = np.expand_dims(img, axis=2)
+ return img
+
+
+def single2tensor5(img):
+ return torch.from_numpy(np.ascontiguousarray(img)).permute(2, 0, 1, 3).float().unsqueeze(0)
+
+
+def single32tensor5(img):
+ return torch.from_numpy(np.ascontiguousarray(img)).float().unsqueeze(0).unsqueeze(0)
+
+
+def single42tensor4(img):
+ return torch.from_numpy(np.ascontiguousarray(img)).permute(2, 0, 1, 3).float()
+
+
+# from skimage.io import imread, imsave
+def tensor2img(tensor, out_type=np.uint8, min_max=(0, 1)):
+ '''
+ Converts a torch Tensor into an image Numpy array of BGR channel order
+ Input: 4D(B,(3/1),H,W), 3D(C,H,W), or 2D(H,W), any range, RGB channel order
+ Output: 3D(H,W,C) or 2D(H,W), [0,255], np.uint8 (default)
+ '''
+ tensor = tensor.squeeze().float().cpu().clamp_(*min_max) # squeeze first, then clamp
+ tensor = (tensor - min_max[0]) / (min_max[1] - min_max[0]) # to range [0,1]
+ n_dim = tensor.dim()
+ if n_dim == 4:
+ n_img = len(tensor)
+ img_np = make_grid(tensor, nrow=int(math.sqrt(n_img)), normalize=False).numpy()
+ img_np = np.transpose(img_np[[2, 1, 0], :, :], (1, 2, 0)) # HWC, BGR
+ elif n_dim == 3:
+ img_np = tensor.numpy()
+ img_np = np.transpose(img_np[[2, 1, 0], :, :], (1, 2, 0)) # HWC, BGR
+ elif n_dim == 2:
+ img_np = tensor.numpy()
+ else:
+ raise TypeError(
+ 'Only support 4D, 3D and 2D tensor. But received with dimension: {:d}'.format(n_dim))
+ if out_type == np.uint8:
+ img_np = (img_np * 255.0).round()
+ # Important. Unlike matlab, numpy.unit8() WILL NOT round by default.
+ return img_np.astype(out_type)
+
+
+'''
+# --------------------------------------------
+# Augmentation, flipe and/or rotate
+# --------------------------------------------
+# The following two are enough.
+# (1) augmet_img: numpy image of WxHxC or WxH
+# (2) augment_img_tensor4: tensor image 1xCxWxH
+# --------------------------------------------
+'''
+
+
+def augment_img(img, mode=0):
+ '''Kai Zhang (github: https://github.com/cszn)
+ '''
+ if mode == 0:
+ return img
+ elif mode == 1:
+ return np.flipud(np.rot90(img))
+ elif mode == 2:
+ return np.flipud(img)
+ elif mode == 3:
+ return np.rot90(img, k=3)
+ elif mode == 4:
+ return np.flipud(np.rot90(img, k=2))
+ elif mode == 5:
+ return np.rot90(img)
+ elif mode == 6:
+ return np.rot90(img, k=2)
+ elif mode == 7:
+ return np.flipud(np.rot90(img, k=3))
+
+
+def augment_img_tensor4(img, mode=0):
+ '''Kai Zhang (github: https://github.com/cszn)
+ '''
+ if mode == 0:
+ return img
+ elif mode == 1:
+ return img.rot90(1, [2, 3]).flip([2])
+ elif mode == 2:
+ return img.flip([2])
+ elif mode == 3:
+ return img.rot90(3, [2, 3])
+ elif mode == 4:
+ return img.rot90(2, [2, 3]).flip([2])
+ elif mode == 5:
+ return img.rot90(1, [2, 3])
+ elif mode == 6:
+ return img.rot90(2, [2, 3])
+ elif mode == 7:
+ return img.rot90(3, [2, 3]).flip([2])
+
+
+def augment_img_tensor(img, mode=0):
+ '''Kai Zhang (github: https://github.com/cszn)
+ '''
+ img_size = img.size()
+ img_np = img.data.cpu().numpy()
+ if len(img_size) == 3:
+ img_np = np.transpose(img_np, (1, 2, 0))
+ elif len(img_size) == 4:
+ img_np = np.transpose(img_np, (2, 3, 1, 0))
+ img_np = augment_img(img_np, mode=mode)
+ img_tensor = torch.from_numpy(np.ascontiguousarray(img_np))
+ if len(img_size) == 3:
+ img_tensor = img_tensor.permute(2, 0, 1)
+ elif len(img_size) == 4:
+ img_tensor = img_tensor.permute(3, 2, 0, 1)
+
+ return img_tensor.type_as(img)
+
+
+def augment_img_np3(img, mode=0):
+ if mode == 0:
+ return img
+ elif mode == 1:
+ return img.transpose(1, 0, 2)
+ elif mode == 2:
+ return img[::-1, :, :]
+ elif mode == 3:
+ img = img[::-1, :, :]
+ img = img.transpose(1, 0, 2)
+ return img
+ elif mode == 4:
+ return img[:, ::-1, :]
+ elif mode == 5:
+ img = img[:, ::-1, :]
+ img = img.transpose(1, 0, 2)
+ return img
+ elif mode == 6:
+ img = img[:, ::-1, :]
+ img = img[::-1, :, :]
+ return img
+ elif mode == 7:
+ img = img[:, ::-1, :]
+ img = img[::-1, :, :]
+ img = img.transpose(1, 0, 2)
+ return img
+
+
+def augment_imgs(img_list, hflip=True, rot=True):
+ # horizontal flip OR rotate
+ hflip = hflip and random.random() < 0.5
+ vflip = rot and random.random() < 0.5
+ rot90 = rot and random.random() < 0.5
+
+ def _augment(img):
+ if hflip:
+ img = img[:, ::-1, :]
+ if vflip:
+ img = img[::-1, :, :]
+ if rot90:
+ img = img.transpose(1, 0, 2)
+ return img
+
+ return [_augment(img) for img in img_list]
+
+
+'''
+# --------------------------------------------
+# modcrop and shave
+# --------------------------------------------
+'''
+
+
+def modcrop(img_in, scale):
+ # img_in: Numpy, HWC or HW
+ img = np.copy(img_in)
+ if img.ndim == 2:
+ H, W = img.shape
+ H_r, W_r = H % scale, W % scale
+ img = img[:H - H_r, :W - W_r]
+ elif img.ndim == 3:
+ H, W, C = img.shape
+ H_r, W_r = H % scale, W % scale
+ img = img[:H - H_r, :W - W_r, :]
+ else:
+ raise ValueError('Wrong img ndim: [{:d}].'.format(img.ndim))
+ return img
+
+
+def shave(img_in, border=0):
+ # img_in: Numpy, HWC or HW
+ img = np.copy(img_in)
+ h, w = img.shape[:2]
+ img = img[border:h-border, border:w-border]
+ return img
+
+
+'''
+# --------------------------------------------
+# image processing process on numpy image
+# channel_convert(in_c, tar_type, img_list):
+# rgb2ycbcr(img, only_y=True):
+# bgr2ycbcr(img, only_y=True):
+# ycbcr2rgb(img):
+# --------------------------------------------
+'''
+
+
+def rgb2ycbcr(img, only_y=True):
+ '''same as matlab rgb2ycbcr
+ only_y: only return Y channel
+ Input:
+ uint8, [0, 255]
+ float, [0, 1]
+ '''
+ in_img_type = img.dtype
+ img.astype(np.float32)
+ if in_img_type != np.uint8:
+ img *= 255.
+ # convert
+ if only_y:
+ rlt = np.dot(img, [65.481, 128.553, 24.966]) / 255.0 + 16.0
+ else:
+ rlt = np.matmul(img, [[65.481, -37.797, 112.0], [128.553, -74.203, -93.786],
+ [24.966, 112.0, -18.214]]) / 255.0 + [16, 128, 128]
+ if in_img_type == np.uint8:
+ rlt = rlt.round()
+ else:
+ rlt /= 255.
+ return rlt.astype(in_img_type)
+
+
+def ycbcr2rgb(img):
+ '''same as matlab ycbcr2rgb
+ Input:
+ uint8, [0, 255]
+ float, [0, 1]
+ '''
+ in_img_type = img.dtype
+ img.astype(np.float32)
+ if in_img_type != np.uint8:
+ img *= 255.
+ # convert
+ rlt = np.matmul(img, [[0.00456621, 0.00456621, 0.00456621], [0, -0.00153632, 0.00791071],
+ [0.00625893, -0.00318811, 0]]) * 255.0 + [-222.921, 135.576, -276.836]
+ if in_img_type == np.uint8:
+ rlt = rlt.round()
+ else:
+ rlt /= 255.
+ return rlt.astype(in_img_type)
+
+
+def bgr2ycbcr(img, only_y=True):
+ '''bgr version of rgb2ycbcr
+ only_y: only return Y channel
+ Input:
+ uint8, [0, 255]
+ float, [0, 1]
+ '''
+ in_img_type = img.dtype
+ img.astype(np.float32)
+ if in_img_type != np.uint8:
+ img *= 255.
+ # convert
+ if only_y:
+ rlt = np.dot(img, [24.966, 128.553, 65.481]) / 255.0 + 16.0
+ else:
+ rlt = np.matmul(img, [[24.966, 112.0, -18.214], [128.553, -74.203, -93.786],
+ [65.481, -37.797, 112.0]]) / 255.0 + [16, 128, 128]
+ if in_img_type == np.uint8:
+ rlt = rlt.round()
+ else:
+ rlt /= 255.
+ return rlt.astype(in_img_type)
+
+
+def channel_convert(in_c, tar_type, img_list):
+ # conversion among BGR, gray and y
+ if in_c == 3 and tar_type == 'gray': # BGR to gray
+ gray_list = [cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) for img in img_list]
+ return [np.expand_dims(img, axis=2) for img in gray_list]
+ elif in_c == 3 and tar_type == 'y': # BGR to y
+ y_list = [bgr2ycbcr(img, only_y=True) for img in img_list]
+ return [np.expand_dims(img, axis=2) for img in y_list]
+ elif in_c == 1 and tar_type == 'RGB': # gray/y to BGR
+ return [cv2.cvtColor(img, cv2.COLOR_GRAY2BGR) for img in img_list]
+ else:
+ return img_list
+
+
+'''
+# --------------------------------------------
+# metric, PSNR and SSIM
+# --------------------------------------------
+'''
+
+
+# --------------------------------------------
+# PSNR
+# --------------------------------------------
+def calculate_psnr(img1, img2, border=0):
+ # img1 and img2 have range [0, 255]
+ #img1 = img1.squeeze()
+ #img2 = img2.squeeze()
+ if not img1.shape == img2.shape:
+ raise ValueError('Input images must have the same dimensions.')
+ h, w = img1.shape[:2]
+ img1 = img1[border:h-border, border:w-border]
+ img2 = img2[border:h-border, border:w-border]
+
+ img1 = img1.astype(np.float64)
+ img2 = img2.astype(np.float64)
+ mse = np.mean((img1 - img2)**2)
+ if mse == 0:
+ return float('inf')
+ return 20 * math.log10(255.0 / math.sqrt(mse))
+
+
+# --------------------------------------------
+# SSIM
+# --------------------------------------------
+def calculate_ssim(img1, img2, border=0):
+ '''calculate SSIM
+ the same outputs as MATLAB's
+ img1, img2: [0, 255]
+ '''
+ #img1 = img1.squeeze()
+ #img2 = img2.squeeze()
+ if not img1.shape == img2.shape:
+ raise ValueError('Input images must have the same dimensions.')
+ h, w = img1.shape[:2]
+ img1 = img1[border:h-border, border:w-border]
+ img2 = img2[border:h-border, border:w-border]
+
+ if img1.ndim == 2:
+ return ssim(img1, img2)
+ elif img1.ndim == 3:
+ if img1.shape[2] == 3:
+ ssims = []
+ for i in range(3):
+ ssims.append(ssim(img1[:,:,i], img2[:,:,i]))
+ return np.array(ssims).mean()
+ elif img1.shape[2] == 1:
+ return ssim(np.squeeze(img1), np.squeeze(img2))
+ else:
+ raise ValueError('Wrong input image dimensions.')
+
+
+def ssim(img1, img2):
+ C1 = (0.01 * 255)**2
+ C2 = (0.03 * 255)**2
+
+ img1 = img1.astype(np.float64)
+ img2 = img2.astype(np.float64)
+ kernel = cv2.getGaussianKernel(11, 1.5)
+ window = np.outer(kernel, kernel.transpose())
+
+ mu1 = cv2.filter2D(img1, -1, window)[5:-5, 5:-5] # valid
+ mu2 = cv2.filter2D(img2, -1, window)[5:-5, 5:-5]
+ mu1_sq = mu1**2
+ mu2_sq = mu2**2
+ mu1_mu2 = mu1 * mu2
+ sigma1_sq = cv2.filter2D(img1**2, -1, window)[5:-5, 5:-5] - mu1_sq
+ sigma2_sq = cv2.filter2D(img2**2, -1, window)[5:-5, 5:-5] - mu2_sq
+ sigma12 = cv2.filter2D(img1 * img2, -1, window)[5:-5, 5:-5] - mu1_mu2
+
+ ssim_map = ((2 * mu1_mu2 + C1) * (2 * sigma12 + C2)) / ((mu1_sq + mu2_sq + C1) *
+ (sigma1_sq + sigma2_sq + C2))
+ return ssim_map.mean()
+
+
+'''
+# --------------------------------------------
+# matlab's bicubic imresize (numpy and torch) [0, 1]
+# --------------------------------------------
+'''
+
+
+# matlab 'imresize' function, now only support 'bicubic'
+def cubic(x):
+ absx = torch.abs(x)
+ absx2 = absx**2
+ absx3 = absx**3
+ return (1.5*absx3 - 2.5*absx2 + 1) * ((absx <= 1).type_as(absx)) + \
+ (-0.5*absx3 + 2.5*absx2 - 4*absx + 2) * (((absx > 1)*(absx <= 2)).type_as(absx))
+
+
+def calculate_weights_indices(in_length, out_length, scale, kernel, kernel_width, antialiasing):
+ if (scale < 1) and (antialiasing):
+ # Use a modified kernel to simultaneously interpolate and antialias- larger kernel width
+ kernel_width = kernel_width / scale
+
+ # Output-space coordinates
+ x = torch.linspace(1, out_length, out_length)
+
+ # Input-space coordinates. Calculate the inverse mapping such that 0.5
+ # in output space maps to 0.5 in input space, and 0.5+scale in output
+ # space maps to 1.5 in input space.
+ u = x / scale + 0.5 * (1 - 1 / scale)
+
+ # What is the left-most pixel that can be involved in the computation?
+ left = torch.floor(u - kernel_width / 2)
+
+ # What is the maximum number of pixels that can be involved in the
+ # computation? Note: it's OK to use an extra pixel here; if the
+ # corresponding weights are all zero, it will be eliminated at the end
+ # of this function.
+ P = math.ceil(kernel_width) + 2
+
+ # The indices of the input pixels involved in computing the k-th output
+ # pixel are in row k of the indices matrix.
+ indices = left.view(out_length, 1).expand(out_length, P) + torch.linspace(0, P - 1, P).view(
+ 1, P).expand(out_length, P)
+
+ # The weights used to compute the k-th output pixel are in row k of the
+ # weights matrix.
+ distance_to_center = u.view(out_length, 1).expand(out_length, P) - indices
+ # apply cubic kernel
+ if (scale < 1) and (antialiasing):
+ weights = scale * cubic(distance_to_center * scale)
+ else:
+ weights = cubic(distance_to_center)
+ # Normalize the weights matrix so that each row sums to 1.
+ weights_sum = torch.sum(weights, 1).view(out_length, 1)
+ weights = weights / weights_sum.expand(out_length, P)
+
+ # If a column in weights is all zero, get rid of it. only consider the first and last column.
+ weights_zero_tmp = torch.sum((weights == 0), 0)
+ if not math.isclose(weights_zero_tmp[0], 0, rel_tol=1e-6):
+ indices = indices.narrow(1, 1, P - 2)
+ weights = weights.narrow(1, 1, P - 2)
+ if not math.isclose(weights_zero_tmp[-1], 0, rel_tol=1e-6):
+ indices = indices.narrow(1, 0, P - 2)
+ weights = weights.narrow(1, 0, P - 2)
+ weights = weights.contiguous()
+ indices = indices.contiguous()
+ sym_len_s = -indices.min() + 1
+ sym_len_e = indices.max() - in_length
+ indices = indices + sym_len_s - 1
+ return weights, indices, int(sym_len_s), int(sym_len_e)
+
+
+# --------------------------------------------
+# imresize for tensor image [0, 1]
+# --------------------------------------------
+def imresize(img, scale, antialiasing=True):
+ # Now the scale should be the same for H and W
+ # input: img: pytorch tensor, CHW or HW [0,1]
+ # output: CHW or HW [0,1] w/o round
+ need_squeeze = True if img.dim() == 2 else False
+ if need_squeeze:
+ img.unsqueeze_(0)
+ in_C, in_H, in_W = img.size()
+ out_C, out_H, out_W = in_C, math.ceil(in_H * scale), math.ceil(in_W * scale)
+ kernel_width = 4
+ kernel = 'cubic'
+
+ # Return the desired dimension order for performing the resize. The
+ # strategy is to perform the resize first along the dimension with the
+ # smallest scale factor.
+ # Now we do not support this.
+
+ # get weights and indices
+ weights_H, indices_H, sym_len_Hs, sym_len_He = calculate_weights_indices(
+ in_H, out_H, scale, kernel, kernel_width, antialiasing)
+ weights_W, indices_W, sym_len_Ws, sym_len_We = calculate_weights_indices(
+ in_W, out_W, scale, kernel, kernel_width, antialiasing)
+ # process H dimension
+ # symmetric copying
+ img_aug = torch.FloatTensor(in_C, in_H + sym_len_Hs + sym_len_He, in_W)
+ img_aug.narrow(1, sym_len_Hs, in_H).copy_(img)
+
+ sym_patch = img[:, :sym_len_Hs, :]
+ inv_idx = torch.arange(sym_patch.size(1) - 1, -1, -1).long()
+ sym_patch_inv = sym_patch.index_select(1, inv_idx)
+ img_aug.narrow(1, 0, sym_len_Hs).copy_(sym_patch_inv)
+
+ sym_patch = img[:, -sym_len_He:, :]
+ inv_idx = torch.arange(sym_patch.size(1) - 1, -1, -1).long()
+ sym_patch_inv = sym_patch.index_select(1, inv_idx)
+ img_aug.narrow(1, sym_len_Hs + in_H, sym_len_He).copy_(sym_patch_inv)
+
+ out_1 = torch.FloatTensor(in_C, out_H, in_W)
+ kernel_width = weights_H.size(1)
+ for i in range(out_H):
+ idx = int(indices_H[i][0])
+ for j in range(out_C):
+ out_1[j, i, :] = img_aug[j, idx:idx + kernel_width, :].transpose(0, 1).mv(weights_H[i])
+
+ # process W dimension
+ # symmetric copying
+ out_1_aug = torch.FloatTensor(in_C, out_H, in_W + sym_len_Ws + sym_len_We)
+ out_1_aug.narrow(2, sym_len_Ws, in_W).copy_(out_1)
+
+ sym_patch = out_1[:, :, :sym_len_Ws]
+ inv_idx = torch.arange(sym_patch.size(2) - 1, -1, -1).long()
+ sym_patch_inv = sym_patch.index_select(2, inv_idx)
+ out_1_aug.narrow(2, 0, sym_len_Ws).copy_(sym_patch_inv)
+
+ sym_patch = out_1[:, :, -sym_len_We:]
+ inv_idx = torch.arange(sym_patch.size(2) - 1, -1, -1).long()
+ sym_patch_inv = sym_patch.index_select(2, inv_idx)
+ out_1_aug.narrow(2, sym_len_Ws + in_W, sym_len_We).copy_(sym_patch_inv)
+
+ out_2 = torch.FloatTensor(in_C, out_H, out_W)
+ kernel_width = weights_W.size(1)
+ for i in range(out_W):
+ idx = int(indices_W[i][0])
+ for j in range(out_C):
+ out_2[j, :, i] = out_1_aug[j, :, idx:idx + kernel_width].mv(weights_W[i])
+ if need_squeeze:
+ out_2.squeeze_()
+ return out_2
+
+
+# --------------------------------------------
+# imresize for numpy image [0, 1]
+# --------------------------------------------
+def imresize_np(img, scale, antialiasing=True):
+ # Now the scale should be the same for H and W
+ # input: img: Numpy, HWC or HW [0,1]
+ # output: HWC or HW [0,1] w/o round
+ img = torch.from_numpy(img)
+ need_squeeze = True if img.dim() == 2 else False
+ if need_squeeze:
+ img.unsqueeze_(2)
+
+ in_H, in_W, in_C = img.size()
+ out_C, out_H, out_W = in_C, math.ceil(in_H * scale), math.ceil(in_W * scale)
+ kernel_width = 4
+ kernel = 'cubic'
+
+ # Return the desired dimension order for performing the resize. The
+ # strategy is to perform the resize first along the dimension with the
+ # smallest scale factor.
+ # Now we do not support this.
+
+ # get weights and indices
+ weights_H, indices_H, sym_len_Hs, sym_len_He = calculate_weights_indices(
+ in_H, out_H, scale, kernel, kernel_width, antialiasing)
+ weights_W, indices_W, sym_len_Ws, sym_len_We = calculate_weights_indices(
+ in_W, out_W, scale, kernel, kernel_width, antialiasing)
+ # process H dimension
+ # symmetric copying
+ img_aug = torch.FloatTensor(in_H + sym_len_Hs + sym_len_He, in_W, in_C)
+ img_aug.narrow(0, sym_len_Hs, in_H).copy_(img)
+
+ sym_patch = img[:sym_len_Hs, :, :]
+ inv_idx = torch.arange(sym_patch.size(0) - 1, -1, -1).long()
+ sym_patch_inv = sym_patch.index_select(0, inv_idx)
+ img_aug.narrow(0, 0, sym_len_Hs).copy_(sym_patch_inv)
+
+ sym_patch = img[-sym_len_He:, :, :]
+ inv_idx = torch.arange(sym_patch.size(0) - 1, -1, -1).long()
+ sym_patch_inv = sym_patch.index_select(0, inv_idx)
+ img_aug.narrow(0, sym_len_Hs + in_H, sym_len_He).copy_(sym_patch_inv)
+
+ out_1 = torch.FloatTensor(out_H, in_W, in_C)
+ kernel_width = weights_H.size(1)
+ for i in range(out_H):
+ idx = int(indices_H[i][0])
+ for j in range(out_C):
+ out_1[i, :, j] = img_aug[idx:idx + kernel_width, :, j].transpose(0, 1).mv(weights_H[i])
+
+ # process W dimension
+ # symmetric copying
+ out_1_aug = torch.FloatTensor(out_H, in_W + sym_len_Ws + sym_len_We, in_C)
+ out_1_aug.narrow(1, sym_len_Ws, in_W).copy_(out_1)
+
+ sym_patch = out_1[:, :sym_len_Ws, :]
+ inv_idx = torch.arange(sym_patch.size(1) - 1, -1, -1).long()
+ sym_patch_inv = sym_patch.index_select(1, inv_idx)
+ out_1_aug.narrow(1, 0, sym_len_Ws).copy_(sym_patch_inv)
+
+ sym_patch = out_1[:, -sym_len_We:, :]
+ inv_idx = torch.arange(sym_patch.size(1) - 1, -1, -1).long()
+ sym_patch_inv = sym_patch.index_select(1, inv_idx)
+ out_1_aug.narrow(1, sym_len_Ws + in_W, sym_len_We).copy_(sym_patch_inv)
+
+ out_2 = torch.FloatTensor(out_H, out_W, in_C)
+ kernel_width = weights_W.size(1)
+ for i in range(out_W):
+ idx = int(indices_W[i][0])
+ for j in range(out_C):
+ out_2[:, i, j] = out_1_aug[:, idx:idx + kernel_width, j].mv(weights_W[i])
+ if need_squeeze:
+ out_2.squeeze_()
+
+ return out_2.numpy()
+
+
+if __name__ == '__main__':
+ print('---')
+# img = imread_uint('test.bmp', 3)
+# img = uint2single(img)
+# img_bicubic = imresize_np(img, 1/4)
\ No newline at end of file
diff --git a/lidm/modules/losses/__init__.py b/lidm/modules/losses/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..ef46a16db54ad2fc925f47c54843a11315575f73
--- /dev/null
+++ b/lidm/modules/losses/__init__.py
@@ -0,0 +1,54 @@
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+
+def adopt_weight(weight, global_step, threshold=0, value=0.):
+ if global_step < threshold:
+ weight = value
+ return weight
+
+
+def hinge_d_loss(logits_real, logits_fake):
+ loss_real = torch.mean(F.relu(1. - logits_real))
+ loss_fake = torch.mean(F.relu(1. + logits_fake))
+ d_loss = 0.5 * (loss_real + loss_fake)
+ return d_loss
+
+
+def vanilla_d_loss(logits_real, logits_fake):
+ d_loss = 0.5 * (
+ torch.mean(torch.nn.functional.softplus(-logits_real)) +
+ torch.mean(torch.nn.functional.softplus(logits_fake)))
+ return d_loss
+
+
+def measure_perplexity(predicted_indices, n_embed):
+ # src: https://github.com/karpathy/deep-vector-quantization/blob/main/model.py
+ # eval cluster perplexity. when perplexity == num_embeddings then all clusters are used exactly equally
+ encodings = F.one_hot(predicted_indices, n_embed).float().reshape(-1, n_embed)
+ avg_probs = encodings.mean(0)
+ perplexity = (-(avg_probs * torch.log(avg_probs + 1e-10)).sum()).exp()
+ cluster_use = torch.sum(avg_probs > 0)
+ return perplexity, cluster_use
+
+
+def l1(x, y):
+ return torch.abs(x - y)
+
+
+def l2(x, y):
+ return torch.pow((x - y), 2)
+
+
+def square_dist_loss(x, y):
+ return torch.sum((x - y) ** 2, dim=1, keepdim=True)
+
+
+def weights_init(m):
+ classname = m.__class__.__name__
+ if classname.find('Conv') != -1:
+ nn.init.normal_(m.weight.data, 0.0, 0.02)
+ elif classname.find('BatchNorm') != -1:
+ nn.init.normal_(m.weight.data, 1.0, 0.02)
+ nn.init.constant_(m.bias.data, 0)
\ No newline at end of file
diff --git a/lidm/modules/losses/contperceptual.py b/lidm/modules/losses/contperceptual.py
new file mode 100644
index 0000000000000000000000000000000000000000..f9afbe01c2d1f6c30a10a810f92feadc082a7090
--- /dev/null
+++ b/lidm/modules/losses/contperceptual.py
@@ -0,0 +1,110 @@
+import torch
+import torch.nn as nn
+
+from . import weights_init, hinge_d_loss, vanilla_d_loss
+from .discriminator import LiDARNLayerDiscriminator
+from .lpips import LPIPS
+
+
+class LPIPSWithDiscriminator(nn.Module):
+ def __init__(self, disc_start, logvar_init=0.0, kl_weight=1.0, pixelloss_weight=1.0,
+ disc_num_layers=3, disc_in_channels=3, disc_factor=1.0, disc_weight=1.0,
+ p_weight=1.0, use_actnorm=False, disc_conditional=False,
+ disc_loss="hinge", **kwargs):
+
+ super().__init__()
+ assert disc_loss in ["hinge", "vanilla"]
+ self.kl_weight = kl_weight
+ self.pixel_weight = pixelloss_weight
+ self.perceptual_weight = p_weight
+ if p_weight > 0.:
+ self.perceptual_loss = LPIPS().eval()
+ # output log variance
+ self.logvar = nn.Parameter(torch.ones(size=()) * logvar_init)
+
+ self.discriminator = LiDARNLayerDiscriminator(
+ input_nc=disc_in_channels, n_layers=disc_num_layers, use_actnorm=use_actnorm).apply(weights_init)
+ self.discriminator_iter_start = disc_start
+ self.disc_loss = hinge_d_loss if disc_loss == "hinge" else vanilla_d_loss
+ self.disc_factor = disc_factor
+ self.discriminator_weight = disc_weight
+ self.disc_conditional = disc_conditional
+
+ def calculate_adaptive_weight(self, nll_loss, g_loss, last_layer=None):
+ if last_layer is not None:
+ nll_grads = torch.autograd.grad(nll_loss, last_layer, retain_graph=True)[0]
+ g_grads = torch.autograd.grad(g_loss, last_layer, retain_graph=True)[0]
+ else:
+ nll_grads = torch.autograd.grad(nll_loss, self.last_layer[0], retain_graph=True)[0]
+ g_grads = torch.autograd.grad(g_loss, self.last_layer[0], retain_graph=True)[0]
+
+ d_weight = torch.norm(nll_grads) / (torch.norm(g_grads) + 1e-4)
+ d_weight = torch.clamp(d_weight, 0.0, 1e4).detach()
+ d_weight = d_weight * self.discriminator_weight
+ return d_weight
+
+ def forward(self, inputs, reconstructions, posteriors, optimizer_idx,
+ global_step, last_layer=None, cond=None, split="train",
+ weights=None):
+ rec_loss = torch.abs(inputs.contiguous() - reconstructions.contiguous())
+ if self.perceptual_weight > 0:
+ p_loss = self.perceptual_loss(inputs.contiguous(), reconstructions.contiguous())
+ rec_loss = rec_loss + self.perceptual_weight * p_loss
+
+ nll_loss = rec_loss / torch.exp(self.logvar) + self.logvar
+ weighted_nll_loss = nll_loss
+ if weights is not None:
+ weighted_nll_loss = weights*nll_loss
+ weighted_nll_loss = torch.sum(weighted_nll_loss) / weighted_nll_loss.shape[0]
+ nll_loss = torch.sum(nll_loss) / nll_loss.shape[0]
+ kl_loss = posteriors.kl()
+ kl_loss = torch.sum(kl_loss) / kl_loss.shape[0]
+
+ # now the GAN part
+ disc_factor = 0. if global_step > self.discriminator_iter_start else self.disc_factor
+ if optimizer_idx == 0:
+ # generator update
+ if cond is None:
+ assert not self.disc_conditional
+ logits_fake = self.discriminator(reconstructions.contiguous())
+ else:
+ assert self.disc_conditional
+ logits_fake = self.discriminator(torch.cat((reconstructions.contiguous(), cond), dim=1))
+ g_loss = -torch.mean(logits_fake)
+
+ if self.disc_factor > 0.0:
+ try:
+ d_weight = self.calculate_adaptive_weight(nll_loss, g_loss, last_layer=last_layer)
+ except RuntimeError:
+ assert not self.training
+ d_weight = torch.tensor(0.0)
+ else:
+ d_weight = torch.tensor(0.0)
+
+ loss = weighted_nll_loss + self.kl_weight * kl_loss + disc_factor * d_weight * g_loss
+
+ log = {"{}/total_loss".format(split): loss.clone().detach().mean(), "{}/logvar".format(split): self.logvar.detach(),
+ "{}/kl_loss".format(split): kl_loss.detach().mean(), "{}/nll_loss".format(split): nll_loss.detach().mean(),
+ "{}/rec_loss".format(split): rec_loss.detach().mean(),
+ "{}/d_weight".format(split): d_weight.detach(),
+ "{}/disc_factor".format(split): torch.tensor(disc_factor),
+ "{}/g_loss".format(split): g_loss.detach().mean(),
+ }
+ return loss, log
+
+ if optimizer_idx == 1:
+ # second pass for discriminator update
+ if cond is None:
+ logits_real = self.discriminator(inputs.contiguous().detach())
+ logits_fake = self.discriminator(reconstructions.contiguous().detach())
+ else:
+ logits_real = self.discriminator(torch.cat((inputs.contiguous().detach(), cond), dim=1))
+ logits_fake = self.discriminator(torch.cat((reconstructions.contiguous().detach(), cond), dim=1))
+
+ d_loss = disc_factor * self.disc_loss(logits_real, logits_fake)
+
+ log = {"{}/disc_loss".format(split): d_loss.clone().detach().mean(),
+ "{}/logits_real".format(split): logits_real.detach().mean(),
+ "{}/logits_fake".format(split): logits_fake.detach().mean()
+ }
+ return d_loss, log
diff --git a/lidm/modules/losses/discriminator.py b/lidm/modules/losses/discriminator.py
new file mode 100644
index 0000000000000000000000000000000000000000..64969ce102bc0088ccce2769bf43cf575c0a3968
--- /dev/null
+++ b/lidm/modules/losses/discriminator.py
@@ -0,0 +1,216 @@
+import functools
+import torch.nn as nn
+
+
+from ..basic import ActNorm, CircularConv2d
+
+
+class NLayerDiscriminator(nn.Module):
+ """Defines a PatchGAN discriminator as in Pix2Pix
+ --> see https://github.com/junyanz/pytorch-CycleGAN-and-pix2pix/blob/master/models/networks.py
+ """
+ def __init__(self, input_nc=1, output_nc=1, ndf=64, n_layers=3, use_actnorm=False):
+ """Construct a PatchGAN discriminator
+ Parameters:
+ input_nc (int) -- the number of channels in input images
+ ndf (int) -- the number of filters in the last conv layer
+ n_layers (int) -- the number of conv layers in the discriminator
+ norm_layer -- normalization layer
+ """
+ super(NLayerDiscriminator, self).__init__()
+ if not use_actnorm:
+ norm_layer = nn.BatchNorm2d
+ else:
+ norm_layer = ActNorm
+ if type(norm_layer) == functools.partial: # no need to use bias as BatchNorm2d has affine parameters
+ use_bias = norm_layer.func != nn.BatchNorm2d
+ else:
+ use_bias = norm_layer != nn.BatchNorm2d
+
+ kw = 4
+ padw = 1
+ sequence = [nn.Conv2d(input_nc, ndf, kernel_size=kw, stride=2, padding=padw), nn.LeakyReLU(0.2, True)]
+ nf_mult = 1
+ for n in range(1, n_layers): # gradually increase the number of filters
+ nf_mult_prev = nf_mult
+ nf_mult = min(2 ** n, 8)
+ sequence += [
+ nn.Conv2d(ndf * nf_mult_prev, ndf * nf_mult, kernel_size=kw, stride=2, padding=padw, bias=use_bias),
+ norm_layer(ndf * nf_mult),
+ nn.LeakyReLU(0.2, True)
+ ]
+
+ nf_mult_prev = nf_mult
+ nf_mult = min(2 ** n_layers, 8)
+ sequence += [
+ nn.Conv2d(ndf * nf_mult_prev, ndf * nf_mult, kernel_size=kw, stride=1, padding=padw, bias=use_bias),
+ norm_layer(ndf * nf_mult),
+ nn.LeakyReLU(0.2, True)
+ ]
+
+ sequence += [
+ nn.Conv2d(ndf * nf_mult, output_nc, kernel_size=kw, stride=1, padding=padw)] # output 1 channel prediction map
+ self.main = nn.Sequential(*sequence)
+
+ def forward(self, input):
+ """Standard forward."""
+ return self.main(input)
+
+
+class LiDARNLayerDiscriminator(nn.Module):
+ """Modified PatchGAN discriminator from Pix2Pix
+ --> see https://github.com/junyanz/pytorch-CycleGAN-and-pix2pix/blob/master/models/networks.py
+ """
+ def __init__(self, input_nc=1, output_nc=1, ndf=64, n_layers=3, use_actnorm=False):
+ """Construct a PatchGAN discriminator
+ Parameters:
+ input_nc (int) -- the number of channels in input images
+ ndf (int) -- the number of filters in the last conv layer
+ n_layers (int) -- the number of conv layers in the discriminator
+ norm_layer -- normalization layer
+ """
+ super(LiDARNLayerDiscriminator, self).__init__()
+ if not use_actnorm:
+ norm_layer = nn.BatchNorm2d
+ else:
+ norm_layer = ActNorm
+ if type(norm_layer) == functools.partial: # no need to use bias as BatchNorm2d has affine parameters
+ use_bias = norm_layer.func != nn.BatchNorm2d
+ else:
+ use_bias = norm_layer != nn.BatchNorm2d
+
+ kw = (4, 4)
+ sequence = [CircularConv2d(input_nc, ndf, kernel_size=kw, stride=(1, 2), padding=(1, 2, 1, 2)), nn.LeakyReLU(0.2, True)]
+ nf_mult = 1
+ nf_mult_prev = 1
+ for n in range(1, n_layers): # gradually increase the number of filters
+ nf_mult_prev = nf_mult
+ nf_mult = min(2 ** n, 8)
+ sequence += [
+ CircularConv2d(ndf * nf_mult_prev, ndf * nf_mult, kernel_size=kw, stride=(1, 2), bias=use_bias, padding=(1, 2, 1, 2)),
+ norm_layer(ndf * nf_mult),
+ nn.LeakyReLU(0.2, True)
+ ]
+
+ nf_mult_prev = nf_mult
+ nf_mult = min(2 ** n_layers, 8)
+ sequence += [
+ CircularConv2d(ndf * nf_mult_prev, ndf * nf_mult, kernel_size=kw, stride=1, bias=use_bias, padding=(1, 2, 1, 2)),
+ norm_layer(ndf * nf_mult),
+ nn.LeakyReLU(0.2, True)
+ ]
+
+ sequence += [
+ CircularConv2d(ndf * nf_mult, output_nc, kernel_size=kw, stride=1, padding=(1, 2, 1, 2))] # output 1 channel prediction map
+ self.main = nn.Sequential(*sequence)
+
+ def forward(self, input):
+ """Standard forward."""
+ return self.main(input)
+
+
+class LiDARNLayerDiscriminatorV2(nn.Module):
+ """Modified PatchGAN discriminator from Pix2Pix (larger receptive field)
+ --> see https://github.com/junyanz/pytorch-CycleGAN-and-pix2pix/blob/master/models/networks.py
+ """
+ def __init__(self, input_nc=1, output_nc=1, ndf=64, n_layers=3, use_actnorm=False):
+ """Construct a PatchGAN discriminator
+ Parameters:
+ input_nc (int) -- the number of channels in input images
+ ndf (int) -- the number of filters in the last conv layer
+ n_layers (int) -- the number of conv layers in the discriminator
+ norm_layer -- normalization layer
+ """
+ super(LiDARNLayerDiscriminatorV2, self).__init__()
+ if not use_actnorm:
+ norm_layer = nn.BatchNorm2d
+ else:
+ norm_layer = ActNorm
+ if type(norm_layer) == functools.partial: # no need to use bias as BatchNorm2d has affine parameters
+ use_bias = norm_layer.func != nn.BatchNorm2d
+ else:
+ use_bias = norm_layer != nn.BatchNorm2d
+
+ kw = (4, 4)
+ sequence = [CircularConv2d(input_nc, ndf, kernel_size=kw, stride=(1, 2), padding=(1, 2, 1, 2)), nn.LeakyReLU(0.2, True),
+ CircularConv2d(ndf, ndf, kernel_size=kw, stride=(1, 2), padding=(1, 2, 1, 2)), nn.LeakyReLU(0.2, True)]
+ nf_mult = 1
+ nf_mult_prev = 1
+ for n in range(1, n_layers): # gradually increase the number of filters
+ nf_mult_prev = nf_mult
+ nf_mult = min(2 ** n, 8)
+ sequence += [
+ CircularConv2d(ndf * nf_mult_prev, ndf * nf_mult, kernel_size=kw, stride=(2, 2), bias=use_bias, padding=(1, 2, 1, 2)),
+ norm_layer(ndf * nf_mult),
+ nn.LeakyReLU(0.2, True)
+ ]
+
+ nf_mult_prev = nf_mult
+ nf_mult = min(2 ** n_layers, 8)
+ sequence += [
+ CircularConv2d(ndf * nf_mult_prev, ndf * nf_mult, kernel_size=kw, stride=1, bias=use_bias, padding=(1, 2, 1, 2)),
+ norm_layer(ndf * nf_mult),
+ nn.LeakyReLU(0.2, True)
+ ]
+
+ sequence += [
+ CircularConv2d(ndf * nf_mult, output_nc, kernel_size=kw, stride=1, padding=(1, 2, 1, 2))] # output 1 channel prediction map
+ self.main = nn.Sequential(*sequence)
+
+ def forward(self, input):
+ """Standard forward."""
+ return self.main(input)
+
+
+class LiDARNLayerDiscriminatorV3(nn.Module):
+ """Modified PatchGAN discriminator from Pix2Pix (larger receptive field)
+ --> see https://github.com/junyanz/pytorch-CycleGAN-and-pix2pix/blob/master/models/networks.py
+ """
+ def __init__(self, input_nc=1, output_nc=1, ndf=64, n_layers=3, use_actnorm=False):
+ """Construct a PatchGAN discriminator
+ Parameters:
+ input_nc (int) -- the number of channels in input images
+ ndf (int) -- the number of filters in the last conv layer
+ n_layers (int) -- the number of conv layers in the discriminator
+ norm_layer -- normalization layer
+ """
+ super(LiDARNLayerDiscriminatorV3, self).__init__()
+ if not use_actnorm:
+ norm_layer = nn.BatchNorm2d
+ else:
+ norm_layer = ActNorm
+ if type(norm_layer) == functools.partial: # no need to use bias as BatchNorm2d has affine parameters
+ use_bias = norm_layer.func != nn.BatchNorm2d
+ else:
+ use_bias = norm_layer != nn.BatchNorm2d
+
+ kw = (4, 4)
+ sequence = [CircularConv2d(input_nc, ndf, kernel_size=(1, 4), stride=(1, 1), padding=(1, 2, 1, 2)), nn.LeakyReLU(0.2, True),
+ CircularConv2d(ndf, ndf, kernel_size=kw, stride=(2, 2), padding=(1, 2, 1, 2)), nn.LeakyReLU(0.2, True)]
+ nf_mult = 1
+ nf_mult_prev = 1
+ for n in range(1, n_layers): # gradually increase the number of filters
+ nf_mult_prev = nf_mult
+ nf_mult = min(2 ** n, 8)
+ sequence += [
+ CircularConv2d(ndf * nf_mult_prev, ndf * nf_mult, kernel_size=kw, stride=(2, 2), bias=use_bias, padding=(1, 2, 1, 2)),
+ norm_layer(ndf * nf_mult),
+ nn.LeakyReLU(0.2, True)
+ ]
+
+ nf_mult_prev = nf_mult
+ nf_mult = min(2 ** n_layers, 8)
+ sequence += [
+ CircularConv2d(ndf * nf_mult_prev, ndf * nf_mult, kernel_size=kw, stride=1, bias=use_bias, padding=(1, 2, 1, 2)),
+ norm_layer(ndf * nf_mult),
+ nn.LeakyReLU(0.2, True)
+ ]
+
+ sequence += [
+ CircularConv2d(ndf * nf_mult, output_nc, kernel_size=kw, stride=1, padding=(1, 2, 1, 2))] # output 1 channel prediction map
+ self.main = nn.Sequential(*sequence)
+
+ def forward(self, input):
+ """Standard forward."""
+ import pdb; pdb.set_trace()
+ return self.main(input)
diff --git a/lidm/modules/losses/geometric.py b/lidm/modules/losses/geometric.py
new file mode 100644
index 0000000000000000000000000000000000000000..62cdc1d71da440bed0a3833c723fad3bec2fdc3d
--- /dev/null
+++ b/lidm/modules/losses/geometric.py
@@ -0,0 +1,78 @@
+from functools import partial
+
+import numpy as np
+import torch
+from torch import nn
+import torch.nn.functional as F
+
+
+class GeoConverter(nn.Module):
+ def __init__(self, curve_length=4, bev_only=False, dataset_config=dict()):
+ super().__init__()
+ self.curve_length = curve_length
+ self.coord_dim = 3 if not bev_only else 2
+ self.convert_fn = self.batch_range2bev if bev_only else self.batch_range2xyz
+
+ fov = dataset_config.fov
+ self.fov_up = fov[0] / 180.0 * np.pi # field of view up in rad
+ self.fov_down = fov[1] / 180.0 * np.pi # field of view down in rad
+ self.fov_range = abs(self.fov_down) + abs(self.fov_up) # get field of view total in rad
+ self.depth_scale = dataset_config.depth_scale
+ self.depth_min, self.depth_max = dataset_config.depth_range
+ self.log_scale = dataset_config.log_scale
+ self.size = dataset_config['size']
+ self.register_conversion()
+
+ def register_conversion(self):
+ scan_x, scan_y = np.meshgrid(np.arange(self.size[1]), np.arange(self.size[0]))
+ scan_x = scan_x.astype(np.float64) / self.size[1]
+ scan_y = scan_y.astype(np.float64) / self.size[0]
+
+ yaw = (np.pi * (scan_x * 2 - 1))
+ pitch = ((1.0 - scan_y) * self.fov_range - abs(self.fov_down))
+
+ to_torch = partial(torch.tensor, dtype=torch.float32)
+
+ self.register_buffer('cos_yaw', torch.cos(to_torch(yaw)))
+ self.register_buffer('sin_yaw', torch.sin(to_torch(yaw)))
+ self.register_buffer('cos_pitch', torch.cos(to_torch(pitch)))
+ self.register_buffer('sin_pitch', torch.sin(to_torch(pitch)))
+
+ def batch_range2xyz(self, imgs):
+ batch_depth = (imgs * 0.5 + 0.5) * self.depth_scale
+ if self.log_scale:
+ batch_depth = torch.exp2(batch_depth) - 1
+ batch_depth = batch_depth.clamp(self.depth_min, self.depth_max)
+
+ batch_x = self.cos_yaw * self.cos_pitch * batch_depth
+ batch_y = -self.sin_yaw * self.cos_pitch * batch_depth
+ batch_z = self.sin_pitch * batch_depth
+ batch_xyz = torch.cat([batch_x, batch_y, batch_z], dim=1)
+
+ return batch_xyz
+
+ def batch_range2bev(self, imgs):
+ batch_depth = (imgs * 0.5 + 0.5) * self.depth_scale
+ if self.log_scale:
+ batch_depth = torch.exp2(batch_depth) - 1
+ batch_depth = batch_depth.clamp(self.depth_min, self.depth_max)
+
+ batch_x = self.cos_yaw * self.cos_pitch * batch_depth
+ batch_y = -self.sin_yaw * self.cos_pitch * batch_depth
+ batch_bev = torch.cat([batch_x, batch_y], dim=1)
+
+ return batch_bev
+
+ def curve_compress(self, batch_coord):
+ compressed_batch_coord = F.avg_pool2d(batch_coord, (1, self.curve_length))
+
+ return compressed_batch_coord
+
+ def forward(self, input):
+ input = input / 2. + .5 # [-1, 1] -> [0, 1]
+
+ input_coord = self.convert_fn(input)
+ if self.curve_length > 1:
+ input_coord = self.curve_compress(input_coord)
+
+ return input_coord
diff --git a/lidm/modules/losses/perceptual.py b/lidm/modules/losses/perceptual.py
new file mode 100644
index 0000000000000000000000000000000000000000..c1255a584194c6fdeae16b5e4df8ecf63e5224ed
--- /dev/null
+++ b/lidm/modules/losses/perceptual.py
@@ -0,0 +1,123 @@
+import hashlib
+import os
+
+import requests
+import torch
+import torch.nn as nn
+
+from tqdm import tqdm
+
+from . import l1, l2
+from ...utils.model_utils import build_model
+
+URL_MAP = {
+}
+
+CKPT_MAP = {
+}
+
+MD5_MAP = {
+}
+
+PERCEPTUAL_TYPE = {
+ 'rangenet_full': [('enc_0', 32), ('enc_1', 64), ('enc_2', 128), ('enc_3', 256), ('enc_4', 512), ('enc_5', 1024),
+ ('dec_4', 512), ('dec_3', 256), ('dec_2', 128), ('dec_1', 64), ('dec_0', 32)],
+ 'rangenet_enc': [('enc_0', 32), ('enc_1', 64), ('enc_2', 128), ('enc_3', 256), ('enc_4', 512), ('enc_5', 1024)],
+ 'rangenet_dec': [('dec_4', 512), ('dec_3', 256), ('dec_2', 128), ('dec_1', 64), ('dec_0', 32)],
+ 'rangenet_final': [('dec_0', 32)]
+}
+
+
+def download(url, local_path, chunk_size=1024):
+ os.makedirs(os.path.split(local_path)[0], exist_ok=True)
+ with requests.get(url, stream=True) as r:
+ total_size = int(r.headers.get("content-length", 0))
+ with tqdm(total=total_size, unit="B", unit_scale=True) as pbar:
+ with open(local_path, "wb") as f:
+ for data in r.iter_content(chunk_size=chunk_size):
+ if data:
+ f.write(data)
+ pbar.update(chunk_size)
+
+
+def md5_hash(path):
+ with open(path, "rb") as f:
+ content = f.read()
+ return hashlib.md5(content).hexdigest()
+
+
+def get_ckpt_path(name, root, check=False):
+ assert name in URL_MAP
+ path = os.path.join(root, CKPT_MAP[name])
+ if not os.path.exists(path) or (check and not md5_hash(path) == MD5_MAP[name]):
+ print("Downloading {} model from {} to {}".format(name, URL_MAP[name], path))
+ download(URL_MAP[name], path)
+ md5 = md5_hash(path)
+ assert md5 == MD5_MAP[name], md5
+ return path
+
+
+class NetLinLayer(nn.Module):
+ """ A single linear layer which does a 1x1 conv """
+
+ def __init__(self, chn_in, chn_out=1, use_dropout=False):
+ super(NetLinLayer, self).__init__()
+ layers = [nn.Dropout(), ] if (use_dropout) else []
+ layers += [nn.Conv2d(chn_in, chn_out, 1, stride=1, padding=0, bias=False), ]
+ self.model = nn.Sequential(*layers)
+
+
+class PerceptualLoss(nn.Module):
+ def __init__(self, ptype, depth_scale, log_scale=True, use_dropout=True, lpips=False, p_loss='l1'):
+ super().__init__()
+ self.depth_scale = depth_scale
+ self.log_scale = log_scale
+
+ if p_loss == "l1":
+ self.p_loss = l1
+ else:
+ self.p_loss = l2
+
+ self.chns = PERCEPTUAL_TYPE[ptype]
+ self.return_list = [name for name, _ in self.chns]
+ self.loss_scale = [5.0, 3.39, 2.29, 1.61, 0.895] # predefined based on the loss of each stage after a few epochs (refer )
+ self.net = build_model('kitti', 'rangenet')
+ self.lin_list = nn.ModuleList([NetLinLayer(ch, use_dropout=use_dropout) for _, ch in self.chns]) if lpips else None
+ for param in self.parameters():
+ param.requires_grad = False
+
+ @staticmethod
+ def normalize_tensor(x, eps=1e-10):
+ norm_factor = torch.sqrt(torch.sum(x ** 2, dim=1, keepdim=True))
+ return x / (norm_factor + eps)
+
+ @staticmethod
+ def spatial_average(x, keepdim=True):
+ return x.mean([2, 3], keepdim=keepdim)
+
+ def preprocess(self, *inputs):
+ assert len(inputs) == 2, 'input with both depth images and coord images'
+ depth_img, xyz_img = inputs
+
+ # scale to standard rangenet input
+ depth_img = (depth_img * 0.5 + 0.5) * self.depth_scale
+ if self.log_scale:
+ depth_img = torch.exp2(depth_img) - 1
+
+ img = torch.cat([depth_img, xyz_img], 1)
+ return img
+
+ def forward(self, target, input):
+ in0_input, in1_input = self.preprocess(*input), self.preprocess(*target)
+ outs0, outs1 = self.net(in0_input, return_list=self.return_list), self.net(in1_input, return_list=self.return_list)
+
+ val_list = []
+ for i, (name, _) in enumerate(self.chns):
+ feats0, feats1 = self.normalize_tensor(outs0[name].to(in0_input.device)), \
+ self.normalize_tensor(outs1[name].to(in0_input.device))
+ diffs = self.p_loss(feats0, feats1)
+ res = self.lin_list[i].model(diffs) if self.lin_list is not None else diffs.mean(1, keepdim=True)
+ res = self.spatial_average(res, keepdim=True) * self.loss_scale[i]
+ val_list.append(res)
+ val = sum(val_list)
+ return val
diff --git a/lidm/modules/losses/vqperceptual.py b/lidm/modules/losses/vqperceptual.py
new file mode 100644
index 0000000000000000000000000000000000000000..e9cd7220338d6beacdf5e70108c6279c1fe18b7c
--- /dev/null
+++ b/lidm/modules/losses/vqperceptual.py
@@ -0,0 +1,176 @@
+import torch
+from torch import nn
+
+from . import weights_init, l1, l2, hinge_d_loss, vanilla_d_loss, measure_perplexity, square_dist_loss
+from .geometric import GeoConverter
+from .discriminator import NLayerDiscriminator, LiDARNLayerDiscriminator, LiDARNLayerDiscriminatorV2
+from .perceptual import PerceptualLoss
+
+VERSION2DISC = {'v0': NLayerDiscriminator, 'v1': LiDARNLayerDiscriminator, 'v2': LiDARNLayerDiscriminatorV2}
+
+
+class VQGeoLPIPSWithDiscriminator(nn.Module):
+ def __init__(self, disc_start, codebook_weight=1.0, pixelloss_weight=1.0,
+ disc_num_layers=3, disc_in_channels=3, disc_out_channels=1, disc_factor=1.0, disc_weight=1.0,
+ mask_factor=0.0, use_actnorm=False, disc_conditional=False,
+ disc_ndf=64, disc_loss="hinge", n_classes=None, pixel_loss="l1", disc_version='v1',
+ geo_factor=1.0, curve_length=4, perceptual_factor=1.0, perceptual_type='rangenet_final',
+ dataset_config=dict()):
+ super().__init__()
+ assert disc_loss in ["hinge", "vanilla"]
+ assert pixel_loss in ["l1", "l2"]
+ self.codebook_weight = codebook_weight
+ self.pixel_weight = pixelloss_weight
+ self.mask_factor = mask_factor
+ self.geo_factor = geo_factor
+
+ # scale of reconstruction loss
+ self.rec_scale = 1
+ if mask_factor > 0:
+ self.rec_scale += 1.
+ if geo_factor > 0:
+ self.rec_scale += 1.
+ if perceptual_factor > 0:
+ self.rec_scale += 1.
+
+ if pixel_loss == "l1":
+ self.pixel_loss = l1
+ else:
+ self.pixel_loss = l2
+
+ self.perceptual_factor = perceptual_factor
+ if perceptual_factor > 0.:
+ print(f"{self.__class__.__name__}: Running with LPIPS.")
+ self.perceptual_loss = PerceptualLoss(perceptual_type, dataset_config.depth_scale,
+ dataset_config.log_scale).eval()
+
+ disc_cls = VERSION2DISC[disc_version]
+ self.discriminator = disc_cls(input_nc=disc_in_channels,
+ output_nc=disc_out_channels,
+ n_layers=disc_num_layers,
+ use_actnorm=use_actnorm,
+ ndf=disc_ndf).apply(weights_init)
+ self.discriminator_iter_start = disc_start
+ if disc_loss == "hinge":
+ self.disc_loss = hinge_d_loss
+ elif disc_loss == "vanilla":
+ self.disc_loss = vanilla_d_loss
+ else:
+ raise ValueError(f"Unknown GAN loss '{disc_loss}'.")
+ print(f"VQGeoLPIPSWithDiscriminator running with {disc_loss} loss.")
+ self.disc_factor = disc_factor
+ self.discriminator_weight = disc_weight
+ self.disc_conditional = disc_conditional
+ self.n_classes = n_classes
+
+ self.geometry_converter = GeoConverter(curve_length, False, dataset_config) # force converting xyz output
+ self.geo_loss = square_dist_loss
+
+ def calculate_adaptive_weight(self, nll_loss, g_loss, last_layer=None):
+ if last_layer is not None:
+ nll_grads = torch.autograd.grad(nll_loss, last_layer, retain_graph=True)[0]
+ g_grads = torch.autograd.grad(g_loss, last_layer, retain_graph=True)[0]
+ else:
+ nll_grads = torch.autograd.grad(nll_loss, self.last_layer[0], retain_graph=True)[0]
+ g_grads = torch.autograd.grad(g_loss, self.last_layer[0], retain_graph=True)[0]
+
+ d_weight = torch.norm(nll_grads) / (torch.norm(g_grads) + 1e-4)
+ d_weight = torch.clamp(d_weight, 0.0, 1e4).detach()
+ d_weight = d_weight * self.discriminator_weight
+ return d_weight
+
+ def forward(self, codebook_loss, inputs, reconstructions, optimizer_idx,
+ global_step, last_layer=None, cond=None, split="train", predicted_indices=None, masks=None):
+ input_coord = self.geometry_converter(inputs)
+ rec_coord = self.geometry_converter(reconstructions[:, 0:1].contiguous())
+
+ ############# Reconstruction #############
+ # pixel reconstruction loss
+ if self.mask_factor > 0 and masks is not None:
+ pixel_rec_loss = self.pixel_loss(inputs.contiguous(), reconstructions[:, 0:1].contiguous())
+ mask_rec_loss = self.pixel_loss(masks.contiguous(), reconstructions[:, 1:2].contiguous()) * self.mask_factor
+ else:
+ pixel_rec_loss = self.pixel_loss(inputs.contiguous(), reconstructions.contiguous())
+ mask_rec_loss = torch.tensor(0.0)
+
+ # geometry reconstruction loss (bev)
+ if self.geo_factor > 0:
+ geo_rec_loss = self.geo_loss(input_coord[:, :2], rec_coord[:, :2]) * self.geo_factor
+ else:
+ geo_rec_loss = torch.tensor(0.0)
+
+ # perceptual loss
+ if self.perceptual_factor > 0:
+ perceptual_loss = self.perceptual_loss((inputs.contiguous(), input_coord),
+ (reconstructions[:, 0:1].contiguous(), rec_coord)) * self.perceptual_factor
+ else:
+ perceptual_loss = torch.tensor(0.0)
+
+ # overall reconstruction loss
+ rec_loss = (pixel_rec_loss + mask_rec_loss + geo_rec_loss + perceptual_loss) / self.rec_scale
+ nll_loss = rec_loss
+ nll_loss = torch.mean(nll_loss)
+
+ ############# GAN #############
+ disc_factor = 0. if global_step > self.discriminator_iter_start else self.disc_factor
+ # update generator (input: img, mask, coord, [cond])
+ if optimizer_idx == 0:
+ disc_recons = reconstructions.contiguous()
+ if self.geo_factor > 0:
+ disc_recons = torch.cat((disc_recons, rec_coord[:, :2].contiguous()), dim=1)
+ if cond is not None and self.disc_conditional:
+ disc_recons = torch.cat((disc_recons, cond), dim=1)
+ logits_fake = self.discriminator(disc_recons)
+
+ # adversarial loss
+ g_loss = -torch.mean(logits_fake)
+
+ try:
+ d_weight = self.calculate_adaptive_weight(nll_loss, g_loss, last_layer=last_layer)
+ except RuntimeError:
+ assert not self.training
+ d_weight = torch.tensor(0.0)
+
+ loss = nll_loss + d_weight * disc_factor * g_loss + self.codebook_weight * codebook_loss.mean()
+
+ log = {"{}/total_loss".format(split): loss.clone().detach().mean(),
+ "{}/quant_loss".format(split): codebook_loss.detach().mean(),
+ "{}/rec_loss".format(split): rec_loss.detach().mean(),
+ "{}/pix_rec_loss".format(split): pixel_rec_loss.detach().mean(),
+ "{}/geo_rec_loss".format(split): geo_rec_loss.detach().mean(),
+ "{}/mask_rec_loss".format(split): mask_rec_loss.detach().mean(),
+ "{}/perceptual_loss".format(split): perceptual_loss.detach().mean(),
+ "{}/d_weight".format(split): d_weight.detach(),
+ "{}/disc_factor".format(split): torch.tensor(disc_factor),
+ "{}/g_loss".format(split): g_loss.detach().mean()}
+
+ if predicted_indices is not None:
+ assert self.n_classes is not None
+ with torch.no_grad():
+ perplexity, cluster_usage = measure_perplexity(predicted_indices, self.n_classes)
+ log[f"{split}/perplexity"] = perplexity
+ log[f"{split}/cluster_usage"] = cluster_usage
+ return loss, log
+
+ # update discriminator (input: img, mask, coord, [cond])
+ if optimizer_idx == 1:
+ disc_inputs, disc_recons = inputs.contiguous().detach(), reconstructions.contiguous().detach()
+ if self.mask_factor > 0 and masks is not None:
+ disc_inputs = torch.cat((disc_inputs, masks.contiguous().detach()), dim=1)
+ if self.geo_factor > 0:
+ disc_inputs = torch.cat((disc_inputs, input_coord[:, :2].contiguous()), dim=1)
+ disc_recons = torch.cat((disc_recons, rec_coord[:, :2].contiguous()), dim=1)
+ if cond is not None:
+ disc_inputs = torch.cat((disc_inputs, cond), dim=1)
+ disc_recons = torch.cat((disc_recons, cond), dim=1)
+ logits_real = self.discriminator(disc_inputs)
+ logits_fake = self.discriminator(disc_recons)
+
+ # gan loss
+ d_loss = self.disc_loss(logits_real, logits_fake) * disc_factor
+
+ log = {"{}/disc_loss".format(split): d_loss.clone().detach().mean(),
+ "{}/logits_real".format(split): logits_real.detach().mean(),
+ "{}/logits_fake".format(split): logits_fake.detach().mean()}
+
+ return d_loss, log
diff --git a/lidm/modules/minkowskinet/__init__.py b/lidm/modules/minkowskinet/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/lidm/modules/minkowskinet/model.py b/lidm/modules/minkowskinet/model.py
new file mode 100644
index 0000000000000000000000000000000000000000..daa36a5009e96f3331653341c6b185627d688a19
--- /dev/null
+++ b/lidm/modules/minkowskinet/model.py
@@ -0,0 +1,141 @@
+import torch
+import torch.nn as nn
+
+try:
+ import torchsparse
+ import torchsparse.nn as spnn
+ from ..ts import basic_blocks
+except ImportError:
+ raise Exception('Required ts lib. Reference: https://github.com/mit-han-lab/torchsparse/tree/v1.4.0')
+
+
+class Model(nn.Module):
+ def __init__(self, config):
+ super().__init__()
+
+ cr = config.model_params.cr
+ cs = config.model_params.layer_num
+ cs = [int(cr * x) for x in cs]
+
+ self.pres = self.vres = config.model_params.voxel_size
+ self.num_classes = config.model_params.num_class
+
+ self.stem = nn.Sequential(
+ spnn.Conv3d(config.model_params.input_dims, cs[0], kernel_size=3, stride=1),
+ spnn.BatchNorm(cs[0]), spnn.ReLU(True),
+ spnn.Conv3d(cs[0], cs[0], kernel_size=3, stride=1),
+ spnn.BatchNorm(cs[0]), spnn.ReLU(True))
+
+ self.stage1 = nn.Sequential(
+ basic_blocks.BasicConvolutionBlock(cs[0], cs[0], ks=2, stride=2, dilation=1),
+ basic_blocks.ResidualBlock(cs[0], cs[1], ks=3, stride=1, dilation=1),
+ basic_blocks.ResidualBlock(cs[1], cs[1], ks=3, stride=1, dilation=1),
+ )
+
+ self.stage2 = nn.Sequential(
+ basic_blocks.BasicConvolutionBlock(cs[1], cs[1], ks=2, stride=2, dilation=1),
+ basic_blocks.ResidualBlock(cs[1], cs[2], ks=3, stride=1, dilation=1),
+ basic_blocks.ResidualBlock(cs[2], cs[2], ks=3, stride=1, dilation=1),
+ )
+
+ self.stage3 = nn.Sequential(
+ basic_blocks.BasicConvolutionBlock(cs[2], cs[2], ks=2, stride=2, dilation=1),
+ basic_blocks.ResidualBlock(cs[2], cs[3], ks=3, stride=1, dilation=1),
+ basic_blocks.ResidualBlock(cs[3], cs[3], ks=3, stride=1, dilation=1),
+ )
+
+ self.stage4 = nn.Sequential(
+ basic_blocks.BasicConvolutionBlock(cs[3], cs[3], ks=2, stride=2, dilation=1),
+ basic_blocks.ResidualBlock(cs[3], cs[4], ks=3, stride=1, dilation=1),
+ basic_blocks.ResidualBlock(cs[4], cs[4], ks=3, stride=1, dilation=1),
+ )
+
+ self.up1 = nn.ModuleList([
+ basic_blocks.BasicDeconvolutionBlock(cs[4], cs[5], ks=2, stride=2),
+ nn.Sequential(
+ basic_blocks.ResidualBlock(cs[5] + cs[3], cs[5], ks=3, stride=1,
+ dilation=1),
+ basic_blocks.ResidualBlock(cs[5], cs[5], ks=3, stride=1, dilation=1),
+ )
+ ])
+
+ self.up2 = nn.ModuleList([
+ basic_blocks.BasicDeconvolutionBlock(cs[5], cs[6], ks=2, stride=2),
+ nn.Sequential(
+ basic_blocks.ResidualBlock(cs[6] + cs[2], cs[6], ks=3, stride=1,
+ dilation=1),
+ basic_blocks.ResidualBlock(cs[6], cs[6], ks=3, stride=1, dilation=1),
+ )
+ ])
+
+ self.up3 = nn.ModuleList([
+ basic_blocks.BasicDeconvolutionBlock(cs[6], cs[7], ks=2, stride=2),
+ nn.Sequential(
+ basic_blocks.ResidualBlock(cs[7] + cs[1], cs[7], ks=3, stride=1,
+ dilation=1),
+ basic_blocks.ResidualBlock(cs[7], cs[7], ks=3, stride=1, dilation=1),
+ )
+ ])
+
+ self.up4 = nn.ModuleList([
+ basic_blocks.BasicDeconvolutionBlock(cs[7], cs[8], ks=2, stride=2),
+ nn.Sequential(
+ basic_blocks.ResidualBlock(cs[8] + cs[0], cs[8], ks=3, stride=1,
+ dilation=1),
+ basic_blocks.ResidualBlock(cs[8], cs[8], ks=3, stride=1, dilation=1),
+ )
+ ])
+
+ self.classifier = nn.Sequential(nn.Linear(cs[8], self.num_classes))
+
+ self.weight_initialization()
+ self.dropout = nn.Dropout(0.3, True)
+
+ def weight_initialization(self):
+ for m in self.modules():
+ if isinstance(m, nn.BatchNorm1d):
+ nn.init.constant_(m.weight, 1)
+ nn.init.constant_(m.bias, 0)
+
+ def forward(self, data_dict, return_logits=False, return_final_logits=False):
+ x = data_dict['lidar']
+ x.C = x.C.int()
+
+ x0 = self.stem(x)
+ x1 = self.stage1(x0)
+ x2 = self.stage2(x1)
+ x3 = self.stage3(x2)
+ x4 = self.stage4(x3)
+
+ if return_logits:
+ output_dict = dict()
+ output_dict['logits'] = x4.F
+ output_dict['batch_indices'] = x4.C[:, -1]
+ return output_dict
+
+ y1 = self.up1[0](x4)
+ y1 = torchsparse.cat([y1, x3])
+ y1 = self.up1[1](y1)
+
+ y2 = self.up2[0](y1)
+ y2 = torchsparse.cat([y2, x2])
+ y2 = self.up2[1](y2)
+
+ y3 = self.up3[0](y2)
+ y3 = torchsparse.cat([y3, x1])
+ y3 = self.up3[1](y3)
+
+ y4 = self.up4[0](y3)
+ y4 = torchsparse.cat([y4, x0])
+ y4 = self.up4[1](y4)
+ if return_final_logits:
+ output_dict = dict()
+ output_dict['logits'] = y4.F
+ output_dict['coords'] = y4.C[:, :3]
+ output_dict['batch_indices'] = y4.C[:, -1]
+ return output_dict
+
+ output = self.classifier(y4.F)
+ data_dict['output'] = output.F
+
+ return data_dict
diff --git a/lidm/modules/rangenet/__init__.py b/lidm/modules/rangenet/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/lidm/modules/rangenet/model.py b/lidm/modules/rangenet/model.py
new file mode 100644
index 0000000000000000000000000000000000000000..752fae9effd476a6bff255e0063675d1cc2f72e2
--- /dev/null
+++ b/lidm/modules/rangenet/model.py
@@ -0,0 +1,372 @@
+#!/usr/bin/env python3
+# This file is covered by the LICENSE file in the root of this project.
+from collections import OrderedDict
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class BasicBlock(nn.Module):
+ def __init__(self, inplanes, planes, bn_d=0.1):
+ super(BasicBlock, self).__init__()
+ self.conv1 = nn.Conv2d(inplanes, planes[0], kernel_size=1,
+ stride=1, padding=0, bias=False)
+ self.bn1 = nn.BatchNorm2d(planes[0], momentum=bn_d)
+ self.relu1 = nn.LeakyReLU(0.1)
+ self.conv2 = nn.Conv2d(planes[0], planes[1], kernel_size=3,
+ stride=1, padding=1, bias=False)
+ self.bn2 = nn.BatchNorm2d(planes[1], momentum=bn_d)
+ self.relu2 = nn.LeakyReLU(0.1)
+
+ def forward(self, x):
+ residual = x
+
+ out = self.conv1(x)
+ out = self.bn1(out)
+ out = self.relu1(out)
+
+ out = self.conv2(out)
+ out = self.bn2(out)
+ out = self.relu2(out)
+
+ out += residual
+ return out
+
+
+# ******************************************************************************
+
+# number of layers per model
+model_blocks = {
+ 21: [1, 1, 2, 2, 1],
+ 53: [1, 2, 8, 8, 4],
+}
+
+
+class Backbone(nn.Module):
+ """
+ Class for DarknetSeg. Subclasses PyTorch's own "nn" module
+ """
+
+ def __init__(self, params):
+ super(Backbone, self).__init__()
+ self.use_range = params["input_depth"]["range"]
+ self.use_xyz = params["input_depth"]["xyz"]
+ self.use_remission = params["input_depth"]["remission"]
+ self.drop_prob = params["dropout"]
+ self.bn_d = params["bn_d"]
+ self.OS = params["OS"]
+ self.layers = params["extra"]["layers"]
+
+ # input depth calc
+ self.input_depth = 0
+ self.input_idxs = []
+ if self.use_range:
+ self.input_depth += 1
+ self.input_idxs.append(0)
+ if self.use_xyz:
+ self.input_depth += 3
+ self.input_idxs.extend([1, 2, 3])
+ if self.use_remission:
+ self.input_depth += 1
+ self.input_idxs.append(4)
+
+ # stride play
+ self.strides = [2, 2, 2, 2, 2]
+ # check current stride
+ current_os = 1
+ for s in self.strides:
+ current_os *= s
+
+ # make the new stride
+ if self.OS > current_os:
+ print("Can't do OS, ", self.OS,
+ " because it is bigger than original ", current_os)
+ else:
+ # redo strides according to needed stride
+ for i, stride in enumerate(reversed(self.strides), 0):
+ if int(current_os) != self.OS:
+ if stride == 2:
+ current_os /= 2
+ self.strides[-1 - i] = 1
+ if int(current_os) == self.OS:
+ break
+
+ # check that darknet exists
+ assert self.layers in model_blocks.keys()
+
+ # generate layers depending on darknet type
+ self.blocks = model_blocks[self.layers]
+
+ # input layer
+ self.conv1 = nn.Conv2d(self.input_depth, 32, kernel_size=3,
+ stride=1, padding=1, bias=False)
+ self.bn1 = nn.BatchNorm2d(32, momentum=self.bn_d)
+ self.relu1 = nn.LeakyReLU(0.1)
+
+ # encoder
+ self.enc1 = self._make_enc_layer(BasicBlock, [32, 64], self.blocks[0],
+ stride=self.strides[0], bn_d=self.bn_d)
+ self.enc2 = self._make_enc_layer(BasicBlock, [64, 128], self.blocks[1],
+ stride=self.strides[1], bn_d=self.bn_d)
+ self.enc3 = self._make_enc_layer(BasicBlock, [128, 256], self.blocks[2],
+ stride=self.strides[2], bn_d=self.bn_d)
+ self.enc4 = self._make_enc_layer(BasicBlock, [256, 512], self.blocks[3],
+ stride=self.strides[3], bn_d=self.bn_d)
+ self.enc5 = self._make_enc_layer(BasicBlock, [512, 1024], self.blocks[4],
+ stride=self.strides[4], bn_d=self.bn_d)
+
+ # for a bit of fun
+ self.dropout = nn.Dropout2d(self.drop_prob)
+
+ # last channels
+ self.last_channels = 1024
+
+ # make layer useful function
+ def _make_enc_layer(self, block, planes, blocks, stride, bn_d=0.1):
+ layers = []
+
+ # downsample
+ layers.append(("conv", nn.Conv2d(planes[0], planes[1],
+ kernel_size=3,
+ stride=[1, stride], dilation=1,
+ padding=1, bias=False)))
+ layers.append(("bn", nn.BatchNorm2d(planes[1], momentum=bn_d)))
+ layers.append(("relu", nn.LeakyReLU(0.1)))
+
+ # blocks
+ inplanes = planes[1]
+ for i in range(0, blocks):
+ layers.append(("residual_{}".format(i),
+ block(inplanes, planes, bn_d)))
+
+ return nn.Sequential(OrderedDict(layers))
+
+ def run_layer(self, x, layer, skips, os):
+ y = layer(x)
+ if y.shape[2] < x.shape[2] or y.shape[3] < x.shape[3]:
+ skips[os] = x.detach()
+ os *= 2
+ x = y
+ return x, skips, os
+
+ def forward(self, x, return_logits=False, return_list=None):
+ # filter input
+ x = x[:, self.input_idxs]
+
+ # run cnn
+ # store for skip connections
+ skips = {}
+ out_dict = {}
+ os = 1
+
+ # first layer
+ x, skips, os = self.run_layer(x, self.conv1, skips, os)
+ x, skips, os = self.run_layer(x, self.bn1, skips, os)
+ x, skips, os = self.run_layer(x, self.relu1, skips, os)
+ if return_list and 'enc_0' in return_list:
+ out_dict['enc_0'] = x.detach().cpu() # 32, 64, 1024
+
+ # all encoder blocks with intermediate dropouts
+ x, skips, os = self.run_layer(x, self.enc1, skips, os)
+ if return_list and 'enc_1' in return_list:
+ out_dict['enc_1'] = x.detach().cpu() # 64, 64, 512
+ x, skips, os = self.run_layer(x, self.dropout, skips, os)
+
+ x, skips, os = self.run_layer(x, self.enc2, skips, os)
+ if return_list and 'enc_2' in return_list:
+ out_dict['enc_2'] = x.detach().cpu() # 128, 64, 256
+ x, skips, os = self.run_layer(x, self.dropout, skips, os)
+
+ x, skips, os = self.run_layer(x, self.enc3, skips, os)
+ if return_list and 'enc_3' in return_list:
+ out_dict['enc_3'] = x.detach().cpu() # 256, 64, 128
+ x, skips, os = self.run_layer(x, self.dropout, skips, os)
+
+ x, skips, os = self.run_layer(x, self.enc4, skips, os)
+ if return_list and 'enc_4' in return_list:
+ out_dict['enc_4'] = x.detach().cpu() # 512, 64, 64
+ x, skips, os = self.run_layer(x, self.dropout, skips, os)
+
+ x, skips, os = self.run_layer(x, self.enc5, skips, os)
+ if return_list and 'enc_5' in return_list:
+ out_dict['enc_5'] = x.detach().cpu() # 1024, 64, 32
+ if return_logits:
+ return x
+
+ x, skips, os = self.run_layer(x, self.dropout, skips, os)
+
+ if return_list is not None:
+ return x, skips, out_dict
+ return x, skips
+
+ def get_last_depth(self):
+ return self.last_channels
+
+ def get_input_depth(self):
+ return self.input_depth
+
+
+class Decoder(nn.Module):
+ """
+ Class for DarknetSeg. Subclasses PyTorch's own "nn" module
+ """
+
+ def __init__(self, params, OS=32, feature_depth=1024):
+ super(Decoder, self).__init__()
+ self.backbone_OS = OS
+ self.backbone_feature_depth = feature_depth
+ self.drop_prob = params["dropout"]
+ self.bn_d = params["bn_d"]
+ self.index = 0
+
+ # stride play
+ self.strides = [2, 2, 2, 2, 2]
+ # check current stride
+ current_os = 1
+ for s in self.strides:
+ current_os *= s
+ # redo strides according to needed stride
+ for i, stride in enumerate(self.strides):
+ if int(current_os) != self.backbone_OS:
+ if stride == 2:
+ current_os /= 2
+ self.strides[i] = 1
+ if int(current_os) == self.backbone_OS:
+ break
+
+ # decoder
+ self.dec5 = self._make_dec_layer(BasicBlock,
+ [self.backbone_feature_depth, 512],
+ bn_d=self.bn_d,
+ stride=self.strides[0])
+ self.dec4 = self._make_dec_layer(BasicBlock, [512, 256], bn_d=self.bn_d,
+ stride=self.strides[1])
+ self.dec3 = self._make_dec_layer(BasicBlock, [256, 128], bn_d=self.bn_d,
+ stride=self.strides[2])
+ self.dec2 = self._make_dec_layer(BasicBlock, [128, 64], bn_d=self.bn_d,
+ stride=self.strides[3])
+ self.dec1 = self._make_dec_layer(BasicBlock, [64, 32], bn_d=self.bn_d,
+ stride=self.strides[4])
+
+ # layer list to execute with skips
+ self.layers = [self.dec5, self.dec4, self.dec3, self.dec2, self.dec1]
+
+ # for a bit of fun
+ self.dropout = nn.Dropout2d(self.drop_prob)
+
+ # last channels
+ self.last_channels = 32
+
+ def _make_dec_layer(self, block, planes, bn_d=0.1, stride=2):
+ layers = []
+
+ # downsample
+ if stride == 2:
+ layers.append(("upconv", nn.ConvTranspose2d(planes[0], planes[1],
+ kernel_size=[1, 4], stride=[1, 2],
+ padding=[0, 1])))
+ else:
+ layers.append(("conv", nn.Conv2d(planes[0], planes[1],
+ kernel_size=3, padding=1)))
+ layers.append(("bn", nn.BatchNorm2d(planes[1], momentum=bn_d)))
+ layers.append(("relu", nn.LeakyReLU(0.1)))
+
+ # blocks
+ layers.append(("residual", block(planes[1], planes, bn_d)))
+
+ return nn.Sequential(OrderedDict(layers))
+
+ def run_layer(self, x, layer, skips, os):
+ feats = layer(x) # up
+ if feats.shape[-1] > x.shape[-1]:
+ os //= 2 # match skip
+ feats = feats + skips[os].detach() # add skip
+ x = feats
+ return x, skips, os
+
+ def forward(self, x, skips, return_logits=False, return_list=None):
+ os = self.backbone_OS
+ out_dict = {}
+
+ # run layers
+ x, skips, os = self.run_layer(x, self.dec5, skips, os)
+ if return_list and 'dec_4' in return_list:
+ out_dict['dec_4'] = x.detach().cpu() # 512, 64, 64
+ x, skips, os = self.run_layer(x, self.dec4, skips, os)
+ if return_list and 'dec_3' in return_list:
+ out_dict['dec_3'] = x.detach().cpu() # 256, 64, 128
+ x, skips, os = self.run_layer(x, self.dec3, skips, os)
+ if return_list and 'dec_2' in return_list:
+ out_dict['dec_2'] = x.detach().cpu() # 128, 64, 256
+ x, skips, os = self.run_layer(x, self.dec2, skips, os)
+ if return_list and 'dec_1' in return_list:
+ out_dict['dec_1'] = x.detach().cpu() # 64, 64, 512
+ x, skips, os = self.run_layer(x, self.dec1, skips, os)
+ if return_list and 'dec_0' in return_list:
+ out_dict['dec_0'] = x.detach().cpu() # 32, 64, 1024
+
+ logits = torch.clone(x).detach()
+ x = self.dropout(x)
+
+ if return_logits:
+ return x, logits
+ if return_list is not None:
+ return out_dict
+ return x
+
+ def get_last_depth(self):
+ return self.last_channels
+
+
+class Model(nn.Module):
+ def __init__(self, config):
+ super().__init__()
+ self.config = config
+ self.backbone = Backbone(params=self.config["backbone"])
+ self.decoder = Decoder(params=self.config["decoder"], OS=self.config["backbone"]["OS"],
+ feature_depth=self.backbone.get_last_depth())
+
+ def load_pretrained_weights(self, path):
+ w_dict = torch.load(path + "/backbone",
+ map_location=lambda storage, loc: storage)
+ self.backbone.load_state_dict(w_dict, strict=True)
+ w_dict = torch.load(path + "/segmentation_decoder",
+ map_location=lambda storage, loc: storage)
+ self.decoder.load_state_dict(w_dict, strict=True)
+
+ def forward(self, x, return_logits=False, return_final_logits=False, return_list=None, agg_type='depth'):
+ if return_logits:
+ logits = self.backbone(x, return_logits)
+ logits = F.adaptive_avg_pool2d(logits, (1, 1)).squeeze()
+ logits = torch.clone(logits).detach().cpu().numpy()
+ return logits
+ elif return_list is not None:
+ x, skips, enc_dict = self.backbone(x, return_list=return_list)
+ dec_dict = self.decoder(x, skips, return_list=return_list)
+ out_dict = {**enc_dict, **dec_dict}
+ return out_dict
+ elif return_final_logits:
+ assert agg_type in ['all', 'sector', 'depth']
+ y, skips = self.backbone(x)
+ y, logits = self.decoder(y, skips, True)
+
+ B, C, H, W = logits.shape
+ N = 16
+
+ # avg all
+ if agg_type == 'all':
+ logits = logits.mean([2, 3])
+ # avg in patch
+ elif agg_type == 'sector':
+ logits = logits.view(B, C, H, N, W // N).mean([2, 4]).reshape(B, -1)
+ # avg in row
+ elif agg_type == 'depth':
+ logits = logits.view(B, C, N, H // N, W).mean([3, 4]).reshape(B, -1)
+
+ logits = torch.clone(logits).detach().cpu().numpy()
+ return logits
+ else:
+ y, skips = self.backbone(x)
+ y = self.decoder(y, skips, False)
+ return y
diff --git a/lidm/modules/spvcnn/__init__.py b/lidm/modules/spvcnn/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/lidm/modules/spvcnn/model.py b/lidm/modules/spvcnn/model.py
new file mode 100644
index 0000000000000000000000000000000000000000..dd7793f8e81cef35f331c0c2c70062023999e83a
--- /dev/null
+++ b/lidm/modules/spvcnn/model.py
@@ -0,0 +1,179 @@
+import torch.nn as nn
+
+try:
+ import torchsparse
+ import torchsparse.nn as spnn
+ from torchsparse import PointTensor
+ from ..ts.utils import initial_voxelize, point_to_voxel, voxel_to_point
+ from ..ts import basic_blocks
+except ImportError:
+ raise Exception('Required torchsparse lib. Reference: https://github.com/mit-han-lab/torchsparse/tree/v1.4.0')
+
+
+class Model(nn.Module):
+ def __init__(self, config):
+ super().__init__()
+ cr = config.model_params.cr
+ cs = config.model_params.layer_num
+ cs = [int(cr * x) for x in cs]
+
+ self.pres = self.vres = config.model_params.voxel_size
+ self.num_classes = config.model_params.num_class
+
+ self.stem = nn.Sequential(
+ spnn.Conv3d(config.model_params.input_dims, cs[0], kernel_size=3, stride=1),
+ spnn.BatchNorm(cs[0]), spnn.ReLU(True),
+ spnn.Conv3d(cs[0], cs[0], kernel_size=3, stride=1),
+ spnn.BatchNorm(cs[0]), spnn.ReLU(True))
+
+ self.stage1 = nn.Sequential(
+ basic_blocks.BasicConvolutionBlock(cs[0], cs[0], ks=2, stride=2, dilation=1),
+ basic_blocks.ResidualBlock(cs[0], cs[1], ks=3, stride=1, dilation=1),
+ basic_blocks.ResidualBlock(cs[1], cs[1], ks=3, stride=1, dilation=1),
+ )
+
+ self.stage2 = nn.Sequential(
+ basic_blocks.BasicConvolutionBlock(cs[1], cs[1], ks=2, stride=2, dilation=1),
+ basic_blocks.ResidualBlock(cs[1], cs[2], ks=3, stride=1, dilation=1),
+ basic_blocks.ResidualBlock(cs[2], cs[2], ks=3, stride=1, dilation=1),
+ )
+
+ self.stage3 = nn.Sequential(
+ basic_blocks.BasicConvolutionBlock(cs[2], cs[2], ks=2, stride=2, dilation=1),
+ basic_blocks.ResidualBlock(cs[2], cs[3], ks=3, stride=1, dilation=1),
+ basic_blocks.ResidualBlock(cs[3], cs[3], ks=3, stride=1, dilation=1),
+ )
+
+ self.stage4 = nn.Sequential(
+ basic_blocks.BasicConvolutionBlock(cs[3], cs[3], ks=2, stride=2, dilation=1),
+ basic_blocks.ResidualBlock(cs[3], cs[4], ks=3, stride=1, dilation=1),
+ basic_blocks.ResidualBlock(cs[4], cs[4], ks=3, stride=1, dilation=1),
+ )
+
+ self.up1 = nn.ModuleList([
+ basic_blocks.BasicDeconvolutionBlock(cs[4], cs[5], ks=2, stride=2),
+ nn.Sequential(
+ basic_blocks.ResidualBlock(cs[5] + cs[3], cs[5], ks=3, stride=1,
+ dilation=1),
+ basic_blocks.ResidualBlock(cs[5], cs[5], ks=3, stride=1, dilation=1),
+ )
+ ])
+
+ self.up2 = nn.ModuleList([
+ basic_blocks.BasicDeconvolutionBlock(cs[5], cs[6], ks=2, stride=2),
+ nn.Sequential(
+ basic_blocks.ResidualBlock(cs[6] + cs[2], cs[6], ks=3, stride=1,
+ dilation=1),
+ basic_blocks.ResidualBlock(cs[6], cs[6], ks=3, stride=1, dilation=1),
+ )
+ ])
+
+ self.up3 = nn.ModuleList([
+ basic_blocks.BasicDeconvolutionBlock(cs[6], cs[7], ks=2, stride=2),
+ nn.Sequential(
+ basic_blocks.ResidualBlock(cs[7] + cs[1], cs[7], ks=3, stride=1,
+ dilation=1),
+ basic_blocks.ResidualBlock(cs[7], cs[7], ks=3, stride=1, dilation=1),
+ )
+ ])
+
+ self.up4 = nn.ModuleList([
+ basic_blocks.BasicDeconvolutionBlock(cs[7], cs[8], ks=2, stride=2),
+ nn.Sequential(
+ basic_blocks.ResidualBlock(cs[8] + cs[0], cs[8], ks=3, stride=1,
+ dilation=1),
+ basic_blocks.ResidualBlock(cs[8], cs[8], ks=3, stride=1, dilation=1),
+ )
+ ])
+
+ self.classifier = nn.Sequential(nn.Linear(cs[8], self.num_classes))
+
+ self.point_transforms = nn.ModuleList([
+ nn.Sequential(
+ nn.Linear(cs[0], cs[4]),
+ nn.BatchNorm1d(cs[4]),
+ nn.ReLU(True),
+ ),
+ nn.Sequential(
+ nn.Linear(cs[4], cs[6]),
+ nn.BatchNorm1d(cs[6]),
+ nn.ReLU(True),
+ ),
+ nn.Sequential(
+ nn.Linear(cs[6], cs[8]),
+ nn.BatchNorm1d(cs[8]),
+ nn.ReLU(True),
+ )
+ ])
+
+ self.weight_initialization()
+ self.dropout = nn.Dropout(0.3, True)
+
+ def weight_initialization(self):
+ for m in self.modules():
+ if isinstance(m, nn.BatchNorm1d):
+ nn.init.constant_(m.weight, 1)
+ nn.init.constant_(m.bias, 0)
+
+ def forward(self, data_dict, return_logits=False, return_final_logits=False):
+ x = data_dict['lidar']
+
+ # x: SparseTensor z: PointTensor
+ z = PointTensor(x.F, x.C.float())
+
+ x0 = initial_voxelize(z, self.pres, self.vres)
+
+ x0 = self.stem(x0)
+ z0 = voxel_to_point(x0, z, nearest=False)
+ z0.F = z0.F
+
+ x1 = point_to_voxel(x0, z0)
+ x1 = self.stage1(x1)
+ x2 = self.stage2(x1)
+ x3 = self.stage3(x2)
+ x4 = self.stage4(x3)
+ z1 = voxel_to_point(x4, z0)
+ z1.F = z1.F + self.point_transforms[0](z0.F)
+
+ y1 = point_to_voxel(x4, z1)
+
+ if return_logits:
+ output_dict = dict()
+ output_dict['logits'] = y1.F
+ output_dict['batch_indices'] = y1.C[:, -1]
+ return output_dict
+
+ y1.F = self.dropout(y1.F)
+ y1 = self.up1[0](y1)
+ y1 = torchsparse.cat([y1, x3])
+ y1 = self.up1[1](y1)
+
+ y2 = self.up2[0](y1)
+ y2 = torchsparse.cat([y2, x2])
+ y2 = self.up2[1](y2)
+ z2 = voxel_to_point(y2, z1)
+ z2.F = z2.F + self.point_transforms[1](z1.F)
+
+ y3 = point_to_voxel(y2, z2)
+ y3.F = self.dropout(y3.F)
+ y3 = self.up3[0](y3)
+ y3 = torchsparse.cat([y3, x1])
+ y3 = self.up3[1](y3)
+
+ y4 = self.up4[0](y3)
+ y4 = torchsparse.cat([y4, x0])
+ y4 = self.up4[1](y4)
+ z3 = voxel_to_point(y4, z2)
+ z3.F = z3.F + self.point_transforms[2](z2.F)
+
+ if return_final_logits:
+ output_dict = dict()
+ output_dict['logits'] = z3.F
+ output_dict['coords'] = z3.C[:, :3]
+ output_dict['batch_indices'] = z3.C[:, -1].long()
+ return output_dict
+
+ # output = self.classifier(z3.F)
+ data_dict['logits'] = z3.F
+
+ return data_dict
diff --git a/lidm/modules/ts/__init__.py b/lidm/modules/ts/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/lidm/modules/ts/basic_blocks.py b/lidm/modules/ts/basic_blocks.py
new file mode 100644
index 0000000000000000000000000000000000000000..15720f7ce0a39aa147367cb53b88e04386499a43
--- /dev/null
+++ b/lidm/modules/ts/basic_blocks.py
@@ -0,0 +1,75 @@
+#!/usr/bin/env python
+# encoding: utf-8
+'''
+@author: Xu Yan
+@file: basic_blocks.py
+@time: 2021/4/14 22:53
+'''
+import torch.nn as nn
+import torchsparse.nn as spnn
+
+
+class BasicConvolutionBlock(nn.Module):
+ def __init__(self, inc, outc, ks=3, stride=1, dilation=1):
+ super().__init__()
+ self.net = nn.Sequential(
+ spnn.Conv3d(
+ inc,
+ outc,
+ kernel_size=ks,
+ dilation=dilation,
+ stride=stride), spnn.BatchNorm(outc),
+ spnn.ReLU(True))
+
+ def forward(self, x):
+ out = self.net(x)
+ return out
+
+
+class BasicDeconvolutionBlock(nn.Module):
+ def __init__(self, inc, outc, ks=3, stride=1):
+ super().__init__()
+ self.net = nn.Sequential(
+ spnn.Conv3d(
+ inc,
+ outc,
+ kernel_size=ks,
+ stride=stride,
+ transposed=True),
+ spnn.BatchNorm(outc),
+ spnn.ReLU(True))
+
+ def forward(self, x):
+ return self.net(x)
+
+
+class ResidualBlock(nn.Module):
+ def __init__(self, inc, outc, ks=3, stride=1, dilation=1):
+ super().__init__()
+ self.net = nn.Sequential(
+ spnn.Conv3d(
+ inc,
+ outc,
+ kernel_size=ks,
+ dilation=dilation,
+ stride=stride), spnn.BatchNorm(outc),
+ spnn.ReLU(True),
+ spnn.Conv3d(
+ outc,
+ outc,
+ kernel_size=ks,
+ dilation=dilation,
+ stride=1),
+ spnn.BatchNorm(outc))
+
+ self.downsample = nn.Sequential() if (inc == outc and stride == 1) else \
+ nn.Sequential(
+ spnn.Conv3d(inc, outc, kernel_size=1, dilation=1, stride=stride),
+ spnn.BatchNorm(outc)
+ )
+
+ self.ReLU = spnn.ReLU(True)
+
+ def forward(self, x):
+ out = self.ReLU(self.net(x) + self.downsample(x))
+ return out
diff --git a/lidm/modules/ts/utils.py b/lidm/modules/ts/utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..77e5e055ce584c5b09bee5857166a950c597ee0d
--- /dev/null
+++ b/lidm/modules/ts/utils.py
@@ -0,0 +1,86 @@
+import torch
+import torchsparse.nn.functional as F
+from torchsparse import PointTensor, SparseTensor
+from torchsparse.nn.utils import get_kernel_offsets
+
+__all__ = ['initial_voxelize', 'point_to_voxel', 'voxel_to_point']
+
+
+# z: PointTensor
+# return: SparseTensor
+def initial_voxelize(z, init_res, after_res):
+ new_float_coord = torch.cat([(z.C[:, :3] * init_res) / after_res, z.C[:, -1].view(-1, 1)], 1)
+
+ pc_hash = F.sphash(torch.floor(new_float_coord).int())
+ sparse_hash = torch.unique(pc_hash)
+ idx_query = F.sphashquery(pc_hash, sparse_hash)
+ counts = F.spcount(idx_query.int(), len(sparse_hash))
+
+ inserted_coords = F.spvoxelize(torch.floor(new_float_coord), idx_query, counts)
+ inserted_coords = torch.round(inserted_coords).int()
+ inserted_feat = F.spvoxelize(z.F, idx_query, counts)
+
+ new_tensor = SparseTensor(inserted_feat, inserted_coords, 1)
+ new_tensor.cmaps.setdefault(new_tensor.stride, new_tensor.coords)
+ z.additional_features['idx_query'][1] = idx_query
+ z.additional_features['counts'][1] = counts
+ z.C = new_float_coord
+
+ return new_tensor
+
+
+# x: SparseTensor, z: PointTensor
+# return: SparseTensor
+def point_to_voxel(x, z):
+ if z.additional_features is None or \
+ z.additional_features.get('idx_query') is None or \
+ z.additional_features['idx_query'].get(x.s) is None:
+ pc_hash = F.sphash(
+ torch.cat([torch.floor(z.C[:, :3] / x.s[0]).int() * x.s[0], z.C[:, -1].int().view(-1, 1)], 1))
+ sparse_hash = F.sphash(x.C)
+ idx_query = F.sphashquery(pc_hash, sparse_hash)
+ counts = F.spcount(idx_query.int(), x.C.shape[0])
+ z.additional_features['idx_query'][x.s] = idx_query
+ z.additional_features['counts'][x.s] = counts
+ else:
+ idx_query = z.additional_features['idx_query'][x.s]
+ counts = z.additional_features['counts'][x.s]
+
+ inserted_feat = F.spvoxelize(z.F, idx_query, counts)
+ new_tensor = SparseTensor(inserted_feat, x.C, x.s)
+ new_tensor.cmaps = x.cmaps
+ new_tensor.kmaps = x.kmaps
+
+ return new_tensor
+
+
+# x: SparseTensor, z: PointTensor
+# return: PointTensor
+def voxel_to_point(x, z, nearest=False):
+ if z.idx_query is None or z.weights is None or z.idx_query.get(x.s) is None or z.weights.get(x.s) is None:
+ off = get_kernel_offsets(2, x.s, 1, device=z.F.device)
+ old_hash = F.sphash(
+ torch.cat([
+ torch.floor(z.C[:, :3] / x.s[0]).int() * x.s[0],
+ z.C[:, -1].int().view(-1, 1)], 1), off)
+ pc_hash = F.sphash(x.C.to(z.F.device))
+ idx_query = F.sphashquery(old_hash, pc_hash)
+ weights = F.calc_ti_weights(z.C, idx_query, scale=x.s[0]).transpose(0, 1).contiguous()
+ idx_query = idx_query.transpose(0, 1).contiguous()
+ if nearest:
+ weights[:, 1:] = 0.
+ idx_query[:, 1:] = -1
+ new_feat = F.spdevoxelize(x.F, idx_query, weights)
+ new_tensor = PointTensor(new_feat, z.C, idx_query=z.idx_query, weights=z.weights)
+ new_tensor.additional_features = z.additional_features
+ new_tensor.idx_query[x.s] = idx_query
+ new_tensor.weights[x.s] = weights
+ z.idx_query[x.s] = idx_query
+ z.weights[x.s] = weights
+
+ else:
+ new_feat = F.spdevoxelize(x.F, z.idx_query.get(x.s), z.weights.get(x.s))
+ new_tensor = PointTensor(new_feat, z.C, idx_query=z.idx_query, weights=z.weights)
+ new_tensor.additional_features = z.additional_features
+
+ return new_tensor
\ No newline at end of file
diff --git a/lidm/modules/x_transformer.py b/lidm/modules/x_transformer.py
new file mode 100644
index 0000000000000000000000000000000000000000..3b969b4e904a4f8ecfcdb4b561ad62c24bff087f
--- /dev/null
+++ b/lidm/modules/x_transformer.py
@@ -0,0 +1,642 @@
+"""shout-out to https://github.com/lucidrains/x-transformers/tree/main/x_transformers"""
+import torch
+from torch import nn, einsum
+import torch.nn.functional as F
+from functools import partial
+from inspect import isfunction
+from collections import namedtuple
+from einops import rearrange, repeat, reduce
+
+# constants
+
+DEFAULT_DIM_HEAD = 64
+
+Intermediates = namedtuple('Intermediates', [
+ 'pre_softmax_attn',
+ 'post_softmax_attn'
+])
+
+LayerIntermediates = namedtuple('Intermediates', [
+ 'hiddens',
+ 'attn_intermediates'
+])
+
+
+class AbsolutePositionalEmbedding(nn.Module):
+ def __init__(self, dim, max_seq_len):
+ super().__init__()
+ self.emb = nn.Embedding(max_seq_len, dim)
+ self.init_()
+
+ def init_(self):
+ nn.init.normal_(self.emb.weight, std=0.02)
+
+ def forward(self, x):
+ n = torch.arange(x.shape[1], device=x.device)
+ return self.emb(n)[None, :, :]
+
+
+class FixedPositionalEmbedding(nn.Module):
+ def __init__(self, dim):
+ super().__init__()
+ inv_freq = 1. / (10000 ** (torch.arange(0, dim, 2).float() / dim))
+ self.register_buffer('inv_freq', inv_freq)
+
+ def forward(self, x, seq_dim=1, offset=0):
+ t = torch.arange(x.shape[seq_dim], device=x.device).type_as(self.inv_freq) + offset
+ sinusoid_inp = torch.einsum('i , j -> i j', t, self.inv_freq)
+ emb = torch.cat((sinusoid_inp.sin(), sinusoid_inp.cos()), dim=-1)
+ return emb[None, :, :]
+
+
+# helpers
+
+def exists(val):
+ return val is not None
+
+
+def default(val, d):
+ if exists(val):
+ return val
+ return d() if isfunction(d) else d
+
+
+def always(val):
+ def inner(*args, **kwargs):
+ return val
+
+ return inner
+
+
+def not_equals(val):
+ def inner(x):
+ return x != val
+
+ return inner
+
+
+def equals(val):
+ def inner(x):
+ return x == val
+
+ return inner
+
+
+def max_neg_value(tensor):
+ return -torch.finfo(tensor.dtype).max
+
+
+# keyword argument helpers
+
+def pick_and_pop(keys, d):
+ values = list(map(lambda key: d.pop(key), keys))
+ return dict(zip(keys, values))
+
+
+def group_dict_by_key(cond, d):
+ return_val = [dict(), dict()]
+ for key in d.keys():
+ match = bool(cond(key))
+ ind = int(not match)
+ return_val[ind][key] = d[key]
+ return (*return_val,)
+
+
+def string_begins_with(prefix, str):
+ return str.startswith(prefix)
+
+
+def group_by_key_prefix(prefix, d):
+ return group_dict_by_key(partial(string_begins_with, prefix), d)
+
+
+def groupby_prefix_and_trim(prefix, d):
+ kwargs_with_prefix, kwargs = group_dict_by_key(partial(string_begins_with, prefix), d)
+ kwargs_without_prefix = dict(map(lambda x: (x[0][len(prefix):], x[1]), tuple(kwargs_with_prefix.items())))
+ return kwargs_without_prefix, kwargs
+
+
+# classes
+class Scale(nn.Module):
+ def __init__(self, value, fn):
+ super().__init__()
+ self.value = value
+ self.fn = fn
+
+ def forward(self, x, **kwargs):
+ x, *rest = self.fn(x, **kwargs)
+ return (x * self.value, *rest)
+
+
+class Rezero(nn.Module):
+ def __init__(self, fn):
+ super().__init__()
+ self.fn = fn
+ self.g = nn.Parameter(torch.zeros(1))
+
+ def forward(self, x, **kwargs):
+ x, *rest = self.fn(x, **kwargs)
+ return (x * self.g, *rest)
+
+
+class ScaleNorm(nn.Module):
+ def __init__(self, dim, eps=1e-5):
+ super().__init__()
+ self.scale = dim ** -0.5
+ self.eps = eps
+ self.g = nn.Parameter(torch.ones(1))
+
+ def forward(self, x):
+ norm = torch.norm(x, dim=-1, keepdim=True) * self.scale
+ return x / norm.clamp(min=self.eps) * self.g
+
+
+class RMSNorm(nn.Module):
+ def __init__(self, dim, eps=1e-8):
+ super().__init__()
+ self.scale = dim ** -0.5
+ self.eps = eps
+ self.g = nn.Parameter(torch.ones(dim))
+
+ def forward(self, x):
+ norm = torch.norm(x, dim=-1, keepdim=True) * self.scale
+ return x / norm.clamp(min=self.eps) * self.g
+
+
+class Residual(nn.Module):
+ def forward(self, x, residual):
+ return x + residual
+
+
+class GRUGating(nn.Module):
+ def __init__(self, dim):
+ super().__init__()
+ self.gru = nn.GRUCell(dim, dim)
+
+ def forward(self, x, residual):
+ gated_output = self.gru(
+ rearrange(x, 'b n d -> (b n) d'),
+ rearrange(residual, 'b n d -> (b n) d')
+ )
+
+ return gated_output.reshape_as(x)
+
+
+# feedforward
+
+class GEGLU(nn.Module):
+ def __init__(self, dim_in, dim_out):
+ super().__init__()
+ self.proj = nn.Linear(dim_in, dim_out * 2)
+
+ def forward(self, x):
+ x, gate = self.proj(x).chunk(2, dim=-1)
+ return x * F.gelu(gate)
+
+
+class FeedForward(nn.Module):
+ def __init__(self, dim, dim_out=None, mult=4, glu=False, dropout=0.):
+ super().__init__()
+ inner_dim = int(dim * mult)
+ dim_out = default(dim_out, dim)
+ project_in = nn.Sequential(
+ nn.Linear(dim, inner_dim),
+ nn.GELU()
+ ) if not glu else GEGLU(dim, inner_dim)
+
+ self.net = nn.Sequential(
+ project_in,
+ nn.Dropout(dropout),
+ nn.Linear(inner_dim, dim_out)
+ )
+
+ def forward(self, x):
+ return self.net(x)
+
+
+# attention.
+class Attention(nn.Module):
+ def __init__(
+ self,
+ dim,
+ dim_head=DEFAULT_DIM_HEAD,
+ heads=8,
+ causal=False,
+ mask=None,
+ talking_heads=False,
+ sparse_topk=None,
+ use_entmax15=False,
+ num_mem_kv=0,
+ dropout=0.,
+ on_attn=False
+ ):
+ super().__init__()
+ if use_entmax15:
+ raise NotImplementedError("Check out entmax activation instead of softmax activation!")
+ self.scale = dim_head ** -0.5
+ self.heads = heads
+ self.causal = causal
+ self.mask = mask
+
+ inner_dim = dim_head * heads
+
+ self.to_q = nn.Linear(dim, inner_dim, bias=False)
+ self.to_k = nn.Linear(dim, inner_dim, bias=False)
+ self.to_v = nn.Linear(dim, inner_dim, bias=False)
+ self.dropout = nn.Dropout(dropout)
+
+ # talking heads
+ self.talking_heads = talking_heads
+ if talking_heads:
+ self.pre_softmax_proj = nn.Parameter(torch.randn(heads, heads))
+ self.post_softmax_proj = nn.Parameter(torch.randn(heads, heads))
+
+ # explicit topk sparse attention
+ self.sparse_topk = sparse_topk
+
+ # entmax
+ # self.attn_fn = entmax15 if use_entmax15 else F.softmax
+ self.attn_fn = F.softmax
+
+ # add memory key / values
+ self.num_mem_kv = num_mem_kv
+ if num_mem_kv > 0:
+ self.mem_k = nn.Parameter(torch.randn(heads, num_mem_kv, dim_head))
+ self.mem_v = nn.Parameter(torch.randn(heads, num_mem_kv, dim_head))
+
+ # attention on attention
+ self.attn_on_attn = on_attn
+ self.to_out = nn.Sequential(nn.Linear(inner_dim, dim * 2), nn.GLU()) if on_attn else nn.Linear(inner_dim, dim)
+
+ def forward(
+ self,
+ x,
+ context=None,
+ mask=None,
+ context_mask=None,
+ rel_pos=None,
+ sinusoidal_emb=None,
+ prev_attn=None,
+ mem=None
+ ):
+ b, n, _, h, talking_heads, device = *x.shape, self.heads, self.talking_heads, x.device
+ kv_input = default(context, x)
+
+ q_input = x
+ k_input = kv_input
+ v_input = kv_input
+
+ if exists(mem):
+ k_input = torch.cat((mem, k_input), dim=-2)
+ v_input = torch.cat((mem, v_input), dim=-2)
+
+ if exists(sinusoidal_emb):
+ # in shortformer, the query would start at a position offset depending on the past cached memory
+ offset = k_input.shape[-2] - q_input.shape[-2]
+ q_input = q_input + sinusoidal_emb(q_input, offset=offset)
+ k_input = k_input + sinusoidal_emb(k_input)
+
+ q = self.to_q(q_input)
+ k = self.to_k(k_input)
+ v = self.to_v(v_input)
+
+ q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h=h), (q, k, v))
+
+ input_mask = None
+ if any(map(exists, (mask, context_mask))):
+ q_mask = default(mask, lambda: torch.ones((b, n), device=device).bool())
+ k_mask = q_mask if not exists(context) else context_mask
+ k_mask = default(k_mask, lambda: torch.ones((b, k.shape[-2]), device=device).bool())
+ q_mask = rearrange(q_mask, 'b i -> b () i ()')
+ k_mask = rearrange(k_mask, 'b j -> b () () j')
+ input_mask = q_mask * k_mask
+
+ if self.num_mem_kv > 0:
+ mem_k, mem_v = map(lambda t: repeat(t, 'h n d -> b h n d', b=b), (self.mem_k, self.mem_v))
+ k = torch.cat((mem_k, k), dim=-2)
+ v = torch.cat((mem_v, v), dim=-2)
+ if exists(input_mask):
+ input_mask = F.pad(input_mask, (self.num_mem_kv, 0), value=True)
+
+ dots = einsum('b h i d, b h j d -> b h i j', q, k) * self.scale
+ mask_value = max_neg_value(dots)
+
+ if exists(prev_attn):
+ dots = dots + prev_attn
+
+ pre_softmax_attn = dots
+
+ if talking_heads:
+ dots = einsum('b h i j, h k -> b k i j', dots, self.pre_softmax_proj).contiguous()
+
+ if exists(rel_pos):
+ dots = rel_pos(dots)
+
+ if exists(input_mask):
+ dots.masked_fill_(~input_mask, mask_value)
+ del input_mask
+
+ if self.causal:
+ i, j = dots.shape[-2:]
+ r = torch.arange(i, device=device)
+ mask = rearrange(r, 'i -> () () i ()') < rearrange(r, 'j -> () () () j')
+ mask = F.pad(mask, (j - i, 0), value=False)
+ dots.masked_fill_(mask, mask_value)
+ del mask
+
+ if exists(self.sparse_topk) and self.sparse_topk < dots.shape[-1]:
+ top, _ = dots.topk(self.sparse_topk, dim=-1)
+ vk = top[..., -1].unsqueeze(-1).expand_as(dots)
+ mask = dots < vk
+ dots.masked_fill_(mask, mask_value)
+ del mask
+
+ attn = self.attn_fn(dots, dim=-1)
+ post_softmax_attn = attn
+
+ attn = self.dropout(attn)
+
+ if talking_heads:
+ attn = einsum('b h i j, h k -> b k i j', attn, self.post_softmax_proj).contiguous()
+
+ out = einsum('b h i j, b h j d -> b h i d', attn, v)
+ out = rearrange(out, 'b h n d -> b n (h d)')
+
+ intermediates = Intermediates(
+ pre_softmax_attn=pre_softmax_attn,
+ post_softmax_attn=post_softmax_attn
+ )
+
+ return self.to_out(out), intermediates
+
+
+class AttentionLayers(nn.Module):
+ def __init__(
+ self,
+ dim,
+ depth,
+ heads=8,
+ causal=False,
+ cross_attend=False,
+ only_cross=False,
+ use_scalenorm=False,
+ use_rmsnorm=False,
+ use_rezero=False,
+ rel_pos_num_buckets=32,
+ rel_pos_max_distance=128,
+ position_infused_attn=False,
+ custom_layers=None,
+ sandwich_coef=None,
+ par_ratio=None,
+ residual_attn=False,
+ cross_residual_attn=False,
+ macaron=False,
+ pre_norm=True,
+ gate_residual=False,
+ **kwargs
+ ):
+ super().__init__()
+ ff_kwargs, kwargs = groupby_prefix_and_trim('ff_', kwargs)
+ attn_kwargs, _ = groupby_prefix_and_trim('attn_', kwargs)
+
+ dim_head = attn_kwargs.get('dim_head', DEFAULT_DIM_HEAD)
+
+ self.dim = dim
+ self.depth = depth
+ self.layers = nn.ModuleList([])
+
+ self.has_pos_emb = position_infused_attn
+ self.pia_pos_emb = FixedPositionalEmbedding(dim) if position_infused_attn else None
+ self.rotary_pos_emb = always(None)
+
+ assert rel_pos_num_buckets <= rel_pos_max_distance, 'number of relative position buckets must be less than the relative position max distance'
+ self.rel_pos = None
+
+ self.pre_norm = pre_norm
+
+ self.residual_attn = residual_attn
+ self.cross_residual_attn = cross_residual_attn
+
+ norm_class = ScaleNorm if use_scalenorm else nn.LayerNorm
+ norm_class = RMSNorm if use_rmsnorm else norm_class
+ norm_fn = partial(norm_class, dim)
+
+ norm_fn = nn.Identity if use_rezero else norm_fn
+ branch_fn = Rezero if use_rezero else None
+
+ if cross_attend and not only_cross:
+ default_block = ('a', 'c', 'f')
+ elif cross_attend and only_cross:
+ default_block = ('c', 'f')
+ else:
+ default_block = ('a', 'f')
+
+ if macaron:
+ default_block = ('f',) + default_block
+
+ if exists(custom_layers):
+ layer_types = custom_layers
+ elif exists(par_ratio):
+ par_depth = depth * len(default_block)
+ assert 1 < par_ratio <= par_depth, 'par ratio out of range'
+ default_block = tuple(filter(not_equals('f'), default_block))
+ par_attn = par_depth // par_ratio
+ depth_cut = par_depth * 2 // 3 # 2 / 3 attention layer cutoff suggested by PAR paper
+ par_width = (depth_cut + depth_cut // par_attn) // par_attn
+ assert len(default_block) <= par_width, 'default block is too large for par_ratio'
+ par_block = default_block + ('f',) * (par_width - len(default_block))
+ par_head = par_block * par_attn
+ layer_types = par_head + ('f',) * (par_depth - len(par_head))
+ elif exists(sandwich_coef):
+ assert sandwich_coef > 0 and sandwich_coef <= depth, 'sandwich coefficient should be less than the depth'
+ layer_types = ('a',) * sandwich_coef + default_block * (depth - sandwich_coef) + ('f',) * sandwich_coef
+ else:
+ layer_types = default_block * depth
+
+ self.layer_types = layer_types
+ self.num_attn_layers = len(list(filter(equals('a'), layer_types)))
+
+ for layer_type in self.layer_types:
+ if layer_type == 'a':
+ layer = Attention(dim, heads=heads, causal=causal, **attn_kwargs)
+ elif layer_type == 'c':
+ layer = Attention(dim, heads=heads, **attn_kwargs)
+ elif layer_type == 'f':
+ layer = FeedForward(dim, **ff_kwargs)
+ layer = layer if not macaron else Scale(0.5, layer)
+ else:
+ raise Exception(f'invalid layer type {layer_type}')
+
+ if isinstance(layer, Attention) and exists(branch_fn):
+ layer = branch_fn(layer)
+
+ if gate_residual:
+ residual_fn = GRUGating(dim)
+ else:
+ residual_fn = Residual()
+
+ self.layers.append(nn.ModuleList([
+ norm_fn(),
+ layer,
+ residual_fn
+ ]))
+
+ def forward(
+ self,
+ x,
+ context=None,
+ mask=None,
+ context_mask=None,
+ mems=None,
+ return_hiddens=False
+ ):
+ hiddens = []
+ intermediates = []
+ prev_attn = None
+ prev_cross_attn = None
+
+ mems = mems.copy() if exists(mems) else [None] * self.num_attn_layers
+
+ for ind, (layer_type, (norm, block, residual_fn)) in enumerate(zip(self.layer_types, self.layers)):
+ is_last = ind == (len(self.layers) - 1)
+
+ if layer_type == 'a':
+ hiddens.append(x)
+ layer_mem = mems.pop(0)
+
+ residual = x
+
+ if self.pre_norm:
+ x = norm(x)
+
+ if layer_type == 'a':
+ out, inter = block(x, mask=mask, sinusoidal_emb=self.pia_pos_emb, rel_pos=self.rel_pos,
+ prev_attn=prev_attn, mem=layer_mem)
+ elif layer_type == 'c':
+ out, inter = block(x, context=context, mask=mask, context_mask=context_mask, prev_attn=prev_cross_attn)
+ elif layer_type == 'f':
+ out = block(x)
+
+ x = residual_fn(out, residual)
+
+ if layer_type in ('a', 'c'):
+ intermediates.append(inter)
+
+ if layer_type == 'a' and self.residual_attn:
+ prev_attn = inter.pre_softmax_attn
+ elif layer_type == 'c' and self.cross_residual_attn:
+ prev_cross_attn = inter.pre_softmax_attn
+
+ if not self.pre_norm and not is_last:
+ x = norm(x)
+
+ if return_hiddens:
+ intermediates = LayerIntermediates(
+ hiddens=hiddens,
+ attn_intermediates=intermediates
+ )
+
+ return x, intermediates
+
+ return x
+
+
+class Encoder(AttentionLayers):
+ def __init__(self, **kwargs):
+ assert 'causal' not in kwargs, 'cannot set causality on encoder'
+ super().__init__(causal=False, **kwargs)
+
+
+class TransformerWrapper(nn.Module):
+ def __init__(
+ self,
+ *,
+ num_tokens,
+ max_seq_len,
+ attn_layers,
+ emb_dim=None,
+ max_mem_len=0.,
+ emb_dropout=0.,
+ num_memory_tokens=None,
+ tie_embedding=False,
+ use_pos_emb=True
+ ):
+ super().__init__()
+ assert isinstance(attn_layers, AttentionLayers), 'attention layers must be one of Encoder or Decoder'
+
+ dim = attn_layers.dim
+ emb_dim = default(emb_dim, dim)
+
+ self.max_seq_len = max_seq_len
+ self.max_mem_len = max_mem_len
+ self.num_tokens = num_tokens
+
+ self.token_emb = nn.Embedding(num_tokens, emb_dim)
+ self.pos_emb = AbsolutePositionalEmbedding(emb_dim, max_seq_len) if (
+ use_pos_emb and not attn_layers.has_pos_emb) else always(0)
+ self.emb_dropout = nn.Dropout(emb_dropout)
+
+ self.project_emb = nn.Linear(emb_dim, dim) if emb_dim != dim else nn.Identity()
+ self.attn_layers = attn_layers
+ self.norm = nn.LayerNorm(dim)
+
+ self.init_()
+
+ self.to_logits = nn.Linear(dim, num_tokens) if not tie_embedding else lambda t: t @ self.token_emb.weight.t()
+
+ # memory tokens (like [cls]) from Memory Transformers paper
+ num_memory_tokens = default(num_memory_tokens, 0)
+ self.num_memory_tokens = num_memory_tokens
+ if num_memory_tokens > 0:
+ self.memory_tokens = nn.Parameter(torch.randn(num_memory_tokens, dim))
+
+ # let funnel encoder know number of memory tokens, if specified
+ if hasattr(attn_layers, 'num_memory_tokens'):
+ attn_layers.num_memory_tokens = num_memory_tokens
+
+ def init_(self):
+ nn.init.normal_(self.token_emb.weight, std=0.02)
+
+ def forward(
+ self,
+ x,
+ return_embeddings=False,
+ mask=None,
+ return_mems=False,
+ return_attn=False,
+ mems=None,
+ **kwargs
+ ):
+ b, n, device, num_mem = *x.shape, x.device, self.num_memory_tokens
+ x = self.token_emb(x)
+ x += self.pos_emb(x)
+ x = self.emb_dropout(x)
+
+ x = self.project_emb(x)
+
+ if num_mem > 0:
+ mem = repeat(self.memory_tokens, 'n d -> b n d', b=b)
+ x = torch.cat((mem, x), dim=1)
+
+ # auto-handle masking after appending memory tokens
+ if exists(mask):
+ mask = F.pad(mask, (num_mem, 0), value=True)
+
+ x, intermediates = self.attn_layers(x, mask=mask, mems=mems, return_hiddens=True, **kwargs)
+ x = self.norm(x)
+
+ mem, x = x[:, :num_mem], x[:, num_mem:]
+
+ out = self.to_logits(x) if not return_embeddings else x
+
+ if return_mems:
+ hiddens = intermediates.hiddens
+ new_mems = list(map(lambda pair: torch.cat(pair, dim=-2), zip(mems, hiddens))) if exists(mems) else hiddens
+ new_mems = list(map(lambda t: t[..., -self.max_mem_len:, :].detach(), new_mems))
+ return out, new_mems
+
+ if return_attn:
+ attn_maps = list(map(lambda t: t.post_softmax_attn, intermediates.attn_intermediates))
+ return out, attn_maps
+
+ return out
diff --git a/lidm/utils/__init__.py b/lidm/utils/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/lidm/utils/aug_utils.py b/lidm/utils/aug_utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..01c130f305afee85e02d2248122059e67ec158ad
--- /dev/null
+++ b/lidm/utils/aug_utils.py
@@ -0,0 +1,107 @@
+import numpy as np
+
+
+def get_lidar_transform(config, split):
+ transform_list = []
+ if config['rotate']:
+ transform_list.append(RandomRotateAligned())
+ if config['flip']:
+ transform_list.append(RandomFlip())
+ return Compose(transform_list) if len(transform_list) > 0 and split == 'train' else None
+
+
+def get_camera_transform(config, split):
+ # import open_clip
+ # transform = open_clip.image_transform((224, 224), split == 'train', resize_longest_max=True)
+ # TODO
+ transform = None
+ return transform
+
+
+def get_anno_transform(config, split):
+ if config['keypoint_drop'] and split == 'train':
+ drop_range = config['keypoint_drop_range'] if 'keypoint_drop_range' in config else (5, 60)
+ transform = RandomKeypointDrop(drop_range)
+ else:
+ transform = None
+ return transform
+
+
+class Compose(object):
+ def __init__(self, transforms):
+ self.transforms = transforms
+
+ def __call__(self, pcd, pcd1=None):
+ for t in self.transforms:
+ pcd, pcd1 = t(pcd, pcd1)
+ return pcd, pcd1
+
+
+class RandomFlip(object):
+ def __init__(self, p=1.):
+ self.p = p
+
+ def __call__(self, coord, coord1=None):
+ if np.random.rand() < self.p:
+ if np.random.rand() < 0.5:
+ coord[:, 0] = -coord[:, 0]
+ if coord1 is not None:
+ coord1[:, 0] = -coord1[:, 0]
+ if np.random.rand() < 0.5:
+ coord[:, 1] = -coord[:, 1]
+ if coord1 is not None:
+ coord1[:, 1] = -coord1[:, 1]
+ return coord, coord1
+
+
+class RandomRotateAligned(object):
+ def __init__(self, rot=np.pi / 4, p=1.):
+ self.rot = rot
+ self.p = p
+
+ def __call__(self, coord, coord1=None):
+ if np.random.rand() < self.p:
+ angle_z = np.random.uniform(-self.rot, self.rot)
+ cos_z, sin_z = np.cos(angle_z), np.sin(angle_z)
+ R = np.array([[cos_z, -sin_z, 0], [sin_z, cos_z, 0], [0, 0, 1]])
+ coord = np.dot(coord, R)
+ if coord1 is not None:
+ coord1 = np.dot(coord1, R)
+ return coord, coord1
+
+
+class RandomKeypointDrop(object):
+ def __init__(self, num_range=(5, 60), p=.5):
+ self.num_range = num_range
+ self.p = p
+
+ def __call__(self, center, category=None):
+ if np.random.rand() < self.p:
+ num = len(center)
+ if num > self.num_range[0]:
+ num_kept = np.random.randint(self.num_range[0], min(self.num_range[1], num))
+ idx_kept = np.random.choice(num, num_kept, replace=False)
+ center, category = center[idx_kept], category[idx_kept]
+ return center, category
+
+
+# class ResizeMaxSize(object):
+# def __init__(self, max_size, interpolation=InterpolationMode.BICUBIC, fn='max', fill=0):
+# super().__init__()
+# if not isinstance(max_size, int):
+# raise TypeError(f"Size should be int. Got {type(max_size)}")
+# self.max_size = max_size
+# self.interpolation = interpolation
+# self.fn = min if fn == 'min' else min
+# self.fill = fill
+#
+# def forward(self, img):
+# width, height = img.size
+# scale = self.max_size / float(max(height, width))
+# if scale != 1.0:
+# new_size = tuple(round(dim * scale) for dim in (height, width))
+# img = F.resize(img, new_size, self.interpolation)
+# pad_h = self.max_size - new_size[0]
+# pad_w = self.max_size - new_size[1]
+# img = F.pad(img, padding=[pad_w//2, pad_h//2, pad_w - pad_w//2, pad_h - pad_h//2], fill=self.fill)
+# return img
diff --git a/lidm/utils/lidar_utils.py b/lidm/utils/lidar_utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..54845771e9fcff03d265a4426f6e214386aee2ba
--- /dev/null
+++ b/lidm/utils/lidar_utils.py
@@ -0,0 +1,206 @@
+import math
+
+import numpy as np
+
+
+def pcd2coord2d(pcd, fov, depth_range, labels=None):
+ # laser parameters
+ fov_up = fov[0] / 180.0 * np.pi # field of view up in rad
+ fov_down = fov[1] / 180.0 * np.pi # field of view down in rad
+ fov_range = abs(fov_down) + abs(fov_up) # get field of view total in rad
+
+ # get depth (distance) of all points
+ depth = np.linalg.norm(pcd, 2, axis=-1)
+
+ # mask points out of range
+ mask = np.logical_and(depth > depth_range[0], depth < depth_range[1])
+ if pcd.ndim == 3:
+ mask = mask.all(axis=1)
+ depth, pcd = depth[mask], pcd[mask]
+
+ # get scan components
+ scan_x, scan_y, scan_z = pcd[..., 0], pcd[..., 1], pcd[..., 2]
+
+ # get angles of all points
+ yaw = -np.arctan2(scan_y, scan_x)
+ pitch = np.arcsin(scan_z / depth)
+
+ # get projections in image coords
+ proj_x = np.clip(0.5 * (yaw / np.pi + 1.0), 0., 1.) # in [0.0, 1.0]
+ proj_y = np.clip(1.0 - (pitch + abs(fov_down)) / fov_range, 0., 1.) # in [0.0, 1.0]
+ proj_coord2d = np.stack([proj_x, proj_y], axis=-1)
+
+ if labels is not None:
+ proj_labels = labels[mask]
+ else:
+ proj_labels = None
+
+ return proj_coord2d, proj_labels
+
+
+def pcd2range(pcd, size, fov, depth_range, remission=None, labels=None, **kwargs):
+ # laser parameters
+ fov_up = fov[0] / 180.0 * np.pi # field of view up in rad
+ fov_down = fov[1] / 180.0 * np.pi # field of view down in rad
+ fov_range = abs(fov_down) + abs(fov_up) # get field of view total in rad
+
+ # get depth (distance) of all points
+ depth = np.linalg.norm(pcd, 2, axis=1)
+
+ # mask points out of range
+ mask = np.logical_and(depth > depth_range[0], depth < depth_range[1])
+ depth, pcd = depth[mask], pcd[mask]
+
+ # get scan components
+ scan_x, scan_y, scan_z = pcd[:, 0], pcd[:, 1], pcd[:, 2]
+
+ # get angles of all points
+ yaw = -np.arctan2(scan_y, scan_x)
+ pitch = np.arcsin(scan_z / depth)
+
+ # get projections in image coords
+ proj_x = 0.5 * (yaw / np.pi + 1.0) # in [0.0, 1.0]
+ proj_y = 1.0 - (pitch + abs(fov_down)) / fov_range # in [0.0, 1.0]
+
+ # scale to image size using angular resolution
+ proj_x *= size[1] # in [0.0, W]
+ proj_y *= size[0] # in [0.0, H]
+
+ # round and clamp for use as index
+ proj_x = np.maximum(0, np.minimum(size[1] - 1, np.floor(proj_x))).astype(np.int32) # in [0,W-1]
+ proj_y = np.maximum(0, np.minimum(size[0] - 1, np.floor(proj_y))).astype(np.int32) # in [0,H-1]
+
+ # order in decreasing depth
+ order = np.argsort(depth)[::-1]
+ proj_x, proj_y = proj_x[order], proj_y[order]
+
+ # project depth
+ depth = depth[order]
+ proj_range = np.full(size, -1, dtype=np.float32)
+ proj_range[proj_y, proj_x] = depth
+
+ # project point feature
+ if remission is not None:
+ remission = remission[mask][order]
+ proj_feature = np.full(size, -1, dtype=np.float32)
+ proj_feature[proj_y, proj_x] = remission
+ elif labels is not None:
+ labels = labels[mask][order]
+ proj_feature = np.full(size, 0, dtype=np.float32)
+ proj_feature[proj_y, proj_x] = labels
+ else:
+ proj_feature = None
+
+ return proj_range, proj_feature
+
+
+def range2pcd(range_img, fov, depth_range, depth_scale, log_scale=True, label=None, color=None, **kwargs):
+ # laser parameters
+ size = range_img.shape
+ fov_up = fov[0] / 180.0 * np.pi # field of view up in rad
+ fov_down = fov[1] / 180.0 * np.pi # field of view down in rad
+ fov_range = abs(fov_down) + abs(fov_up) # get field of view total in rad
+
+ # inverse transform from depth
+ depth = (range_img * depth_scale).flatten()
+ if log_scale:
+ depth = np.exp2(depth) - 1
+
+ scan_x, scan_y = np.meshgrid(np.arange(size[1]), np.arange(size[0]))
+ scan_x = scan_x.astype(np.float64) / size[1]
+ scan_y = scan_y.astype(np.float64) / size[0]
+
+ yaw = (np.pi * (scan_x * 2 - 1)).flatten()
+ pitch = ((1.0 - scan_y) * fov_range - abs(fov_down)).flatten()
+
+ pcd = np.zeros((len(yaw), 3))
+ pcd[:, 0] = np.cos(yaw) * np.cos(pitch) * depth
+ pcd[:, 1] = -np.sin(yaw) * np.cos(pitch) * depth
+ pcd[:, 2] = np.sin(pitch) * depth
+
+ # mask out invalid points
+ mask = np.logical_and(depth > depth_range[0], depth < depth_range[1])
+ pcd = pcd[mask, :]
+
+ # label
+ if label is not None:
+ label = label.flatten()[mask]
+
+ # default point color
+ if color is not None:
+ color = color.reshape(-1, 3)[mask, :]
+ else:
+ color = np.ones((pcd.shape[0], 3)) * [0.7, 0.7, 1]
+
+ return pcd, color, label
+
+
+def range2xyz(range_img, fov, depth_range, depth_scale, log_scale=True, **kwargs):
+ # laser parameters
+ size = range_img.shape
+ fov_up = fov[0] / 180.0 * np.pi # field of view up in rad
+ fov_down = fov[1] / 180.0 * np.pi # field of view down in rad
+ fov_range = abs(fov_down) + abs(fov_up) # get field of view total in rad
+
+ # inverse transform from depth
+ if log_scale:
+ depth = (np.exp2(range_img * depth_scale) - 1)
+ else:
+ depth = range_img
+
+ scan_x, scan_y = np.meshgrid(np.arange(size[1]), np.arange(size[0]))
+ scan_x = scan_x.astype(np.float64) / size[1]
+ scan_y = scan_y.astype(np.float64) / size[0]
+
+ yaw = np.pi * (scan_x * 2 - 1)
+ pitch = (1.0 - scan_y) * fov_range - abs(fov_down)
+
+ xyz = -np.ones((3, *size))
+ xyz[0] = np.cos(yaw) * np.cos(pitch) * depth
+ xyz[1] = -np.sin(yaw) * np.cos(pitch) * depth
+ xyz[2] = np.sin(pitch) * depth
+
+ # mask out invalid points
+ mask = np.logical_and(depth > depth_range[0], depth < depth_range[1])
+ xyz[:, ~mask] = -1
+
+ return xyz
+
+
+def pcd2bev(pcd, x_range, y_range, z_range, resolution, **kwargs):
+ # mask out invalid points
+ mask_x = np.logical_and(pcd[:, 0] > x_range[0], pcd[:, 0] < x_range[1])
+ mask_y = np.logical_and(pcd[:, 1] > y_range[0], pcd[:, 1] < y_range[1])
+ mask_z = np.logical_and(pcd[:, 2] > z_range[0], pcd[:, 2] < z_range[1])
+ mask = mask_x & mask_y & mask_z
+ pcd = pcd[mask]
+
+ # points to bev coords
+ bev_x = np.floor((pcd[:, 0] - x_range[0]) / resolution).astype(np.int32)
+ bev_y = np.floor((pcd[:, 1] - y_range[0]) / resolution).astype(np.int32)
+
+ # 2D bev grid
+ bev_shape = (math.ceil((x_range[1] - x_range[0]) // resolution), math.ceil((y_range[1] - y_range[0]) // resolution))
+ bev_grid = np.zeros(bev_shape, dtype=np.float64)
+
+ # populate the BEV grid with bev coords
+ bev_grid[bev_x, bev_y] = 1
+
+ return bev_grid
+
+
+if __name__ == '__main__':
+ # test = np.loadtxt('test_range.txt')
+ # pcd, _, _ = range2pcd(test, (32, 1024), (10, -30))
+ # np.savetxt('test_pcd.txt', pcd, fmt='%.4f')
+
+ # import matplotlib.pyplot as plt
+ # pcd = np.loadtxt('test_origin.txt')
+ # bev_grid = pcd2bev(pcd)
+ # plt.imshow(bev_grid[:, :, 0], cmap='gray') # Display the BEV for the first height level
+ # plt.savefig('test.png', dpi=300, bbox_inches='tight', pad_inches=0, transparent=True)
+
+ from PIL import Image
+ img = Image.open('assets/kitti/range.png')
+ img.convert('L')
+ img = np.array(img) / 255.
diff --git a/lidm/utils/lr_scheduler.py b/lidm/utils/lr_scheduler.py
new file mode 100644
index 0000000000000000000000000000000000000000..be39da9ca6dacc22bf3df9c7389bbb403a4a3ade
--- /dev/null
+++ b/lidm/utils/lr_scheduler.py
@@ -0,0 +1,98 @@
+import numpy as np
+
+
+class LambdaWarmUpCosineScheduler:
+ """
+ note: use with a base_lr of 1.0
+ """
+ def __init__(self, warm_up_steps, lr_min, lr_max, lr_start, max_decay_steps, verbosity_interval=0):
+ self.lr_warm_up_steps = warm_up_steps
+ self.lr_start = lr_start
+ self.lr_min = lr_min
+ self.lr_max = lr_max
+ self.lr_max_decay_steps = max_decay_steps
+ self.last_lr = 0.
+ self.verbosity_interval = verbosity_interval
+
+ def schedule(self, n, **kwargs):
+ if self.verbosity_interval > 0:
+ if n % self.verbosity_interval == 0: print(f"current step: {n}, recent lr-multiplier: {self.last_lr}")
+ if n < self.lr_warm_up_steps:
+ lr = (self.lr_max - self.lr_start) / self.lr_warm_up_steps * n + self.lr_start
+ self.last_lr = lr
+ return lr
+ else:
+ t = (n - self.lr_warm_up_steps) / (self.lr_max_decay_steps - self.lr_warm_up_steps)
+ t = min(t, 1.0)
+ lr = self.lr_min + 0.5 * (self.lr_max - self.lr_min) * (
+ 1 + np.cos(t * np.pi))
+ self.last_lr = lr
+ return lr
+
+ def __call__(self, n, **kwargs):
+ return self.schedule(n,**kwargs)
+
+
+class LambdaWarmUpCosineScheduler2:
+ """
+ supports repeated iterations, configurable via lists
+ note: use with a base_lr of 1.0.
+ """
+ def __init__(self, warm_up_steps, f_min, f_max, f_start, cycle_lengths, verbosity_interval=0):
+ assert len(warm_up_steps) == len(f_min) == len(f_max) == len(f_start) == len(cycle_lengths)
+ self.lr_warm_up_steps = warm_up_steps
+ self.f_start = f_start
+ self.f_min = f_min
+ self.f_max = f_max
+ self.cycle_lengths = cycle_lengths
+ self.cum_cycles = np.cumsum([0] + list(self.cycle_lengths))
+ self.last_f = 0.
+ self.verbosity_interval = verbosity_interval
+
+ def find_in_interval(self, n):
+ interval = 0
+ for cl in self.cum_cycles[1:]:
+ if n <= cl:
+ return interval
+ interval += 1
+
+ def schedule(self, n, **kwargs):
+ cycle = self.find_in_interval(n)
+ n = n - self.cum_cycles[cycle]
+ if self.verbosity_interval > 0:
+ if n % self.verbosity_interval == 0: print(f"current step: {n}, recent lr-multiplier: {self.last_f}, "
+ f"current cycle {cycle}")
+ if n < self.lr_warm_up_steps[cycle]:
+ f = (self.f_max[cycle] - self.f_start[cycle]) / self.lr_warm_up_steps[cycle] * n + self.f_start[cycle]
+ self.last_f = f
+ return f
+ else:
+ t = (n - self.lr_warm_up_steps[cycle]) / (self.cycle_lengths[cycle] - self.lr_warm_up_steps[cycle])
+ t = min(t, 1.0)
+ f = self.f_min[cycle] + 0.5 * (self.f_max[cycle] - self.f_min[cycle]) * (
+ 1 + np.cos(t * np.pi))
+ self.last_f = f
+ return f
+
+ def __call__(self, n, **kwargs):
+ return self.schedule(n, **kwargs)
+
+
+class LambdaLinearScheduler(LambdaWarmUpCosineScheduler2):
+
+ def schedule(self, n, **kwargs):
+ cycle = self.find_in_interval(n)
+ n = n - self.cum_cycles[cycle]
+ if self.verbosity_interval > 0:
+ if n % self.verbosity_interval == 0: print(f"current step: {n}, recent lr-multiplier: {self.last_f}, "
+ f"current cycle {cycle}")
+
+ if n < self.lr_warm_up_steps[cycle]:
+ f = (self.f_max[cycle] - self.f_start[cycle]) / self.lr_warm_up_steps[cycle] * n + self.f_start[cycle]
+ self.last_f = f
+ return f
+ else:
+ f = self.f_min[cycle] + (self.f_max[cycle] - self.f_min[cycle]) * (self.cycle_lengths[cycle] - n) / (self.cycle_lengths[cycle])
+ self.last_f = f
+ return f
+
diff --git a/lidm/utils/misc_utils.py b/lidm/utils/misc_utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..339e5e27c7a8a5029a9f2dd0cf9f44006b2b756d
--- /dev/null
+++ b/lidm/utils/misc_utils.py
@@ -0,0 +1,243 @@
+import argparse
+import importlib
+import random
+
+import torch
+import numpy as np
+from collections import abc
+from einops import rearrange
+from functools import partial
+
+import multiprocessing as mp
+from threading import Thread
+from queue import Queue
+
+from inspect import isfunction
+from PIL import Image, ImageDraw, ImageFont
+
+
+def set_seed(seed):
+ """
+ Setting of Global Seed for Reproducibility (for inference only)
+
+ refer to: https://pytorch.org/docs/stable/notes/randomness.html
+
+ """
+ np.random.seed(seed)
+ random.seed(seed)
+ torch.manual_seed(seed)
+ torch.cuda.manual_seed(seed)
+
+ torch.backends.cudnn.deterministic = True
+ torch.backends.cudnn.benchmark = False
+
+
+def print_fn(msg, verbose):
+ if verbose:
+ print(msg)
+
+
+def dict2namespace(config):
+ namespace = argparse.Namespace()
+ for key, value in config.items():
+ if isinstance(value, dict):
+ new_value = dict2namespace(value)
+ else:
+ new_value = value
+ setattr(namespace, key, new_value)
+ return namespace
+
+
+def log_txt_as_img(wh, xc, size=10):
+ # wh a tuple of (width, height)
+ # xc a list of captions to plot
+ b = len(xc)
+ txts = list()
+ for bi in range(b):
+ txt = Image.new("RGB", wh, color="white")
+ draw = ImageDraw.Draw(txt)
+ font = ImageFont.truetype('data/DejaVuSans.ttf', size=size)
+ nc = int(40 * (wh[0] / 256))
+ lines = "\n".join(xc[bi][start:start + nc] for start in range(0, len(xc[bi]), nc))
+
+ try:
+ draw.text((0, 0), lines, fill="black", font=font)
+ except UnicodeEncodeError:
+ print("Cant encode string for logging. Skipping.")
+
+ txt = np.array(txt).transpose(2, 0, 1) / 127.5 - 1.0
+ txts.append(txt)
+ txts = np.stack(txts)
+ txts = torch.tensor(txts)
+ return txts
+
+
+def isdepth(x):
+ if not isinstance(x, (torch.Tensor, np.ndarray)):
+ return False
+ return ((len(x.shape) == 4) and (x.shape[1] == 1)) or (len(x.shape) == 3)
+
+
+def ismap(x):
+ if not isinstance(x, (torch.Tensor, np.ndarray)):
+ return False
+ return (len(x.shape) == 4) and (x.shape[1] > 3)
+
+
+def isimage(x):
+ if not isinstance(x, (torch.Tensor, np.ndarray)):
+ return False
+ return (len(x.shape) == 4) and (x.shape[1] == 3 or x.shape[1] == 1)
+
+
+def exists(x):
+ return x is not None
+
+
+def default(val, d):
+ if exists(val):
+ return val
+ return d() if isfunction(d) else d
+
+
+def mean_flat(tensor):
+ """
+ https://github.com/openai/guided-diffusion/blob/27c20a8fab9cb472df5d6bdd6c8d11c8f430b924/guided_diffusion/nn.py#L86
+ Take the mean over all non-batch dimensions.
+ """
+ return tensor.mean(dim=list(range(1, len(tensor.shape))))
+
+
+def count_params(model, verbose=False):
+ total_params = sum(p.numel() for p in model.parameters())
+ if verbose:
+ print(f"{model.__class__.__name__} has {total_params * 1.e-6:.2f} M params.")
+ return total_params
+
+
+def instantiate_from_config(config):
+ if not "target" in config:
+ if config == '__is_first_stage__':
+ return None
+ elif config == "__is_unconditional__":
+ return None
+ raise KeyError("Expected key `target` to instantiate.")
+ return get_obj_from_str(config["target"])(**config.get("params", dict()))
+
+
+def get_obj_from_str(string, reload=False):
+ module, cls = string.rsplit(".", 1)
+ if reload:
+ module_imp = importlib.import_module(module)
+ importlib.reload(module_imp)
+ return getattr(importlib.import_module(module, package=None), cls)
+
+
+def _do_parallel_data_prefetch(func, Q, data, idx, idx_to_fn=False):
+ # create dummy dataset instance
+
+ # run prefetching
+ if idx_to_fn:
+ res = func(data, worker_id=idx)
+ else:
+ res = func(data)
+ Q.put([idx, res])
+ Q.put("Done")
+
+
+def parallel_data_prefetch(
+ func: callable, data, n_proc, target_data_type="ndarray", cpu_intensive=True, use_worker_id=False
+):
+ # if target_data_type not in ["ndarray", "list"]:
+ # raise ValueError(
+ # "Data, which is passed to parallel_data_prefetch has to be either of type list or ndarray."
+ # )
+ if isinstance(data, np.ndarray) and target_data_type == "list":
+ raise ValueError("list expected but function got ndarray.")
+ elif isinstance(data, abc.Iterable):
+ if isinstance(data, dict):
+ print(
+ f'WARNING:"data" argument passed to parallel_data_prefetch is a dict: Using only its values and disregarding keys.'
+ )
+ data = list(data.values())
+ if target_data_type == "ndarray":
+ data = np.asarray(data)
+ else:
+ data = list(data)
+ else:
+ raise TypeError(
+ f"The data, that shall be processed parallel has to be either an np.ndarray or an Iterable, but is actually {type(data)}."
+ )
+
+ if cpu_intensive:
+ Q = mp.Queue(1000)
+ proc = mp.Process
+ else:
+ Q = Queue(1000)
+ proc = Thread
+ # spawn processes
+ if target_data_type == "ndarray":
+ arguments = [
+ [func, Q, part, i, use_worker_id]
+ for i, part in enumerate(np.array_split(data, n_proc))
+ ]
+ else:
+ step = (
+ int(len(data) / n_proc + 1)
+ if len(data) % n_proc != 0
+ else int(len(data) / n_proc)
+ )
+ arguments = [
+ [func, Q, part, i, use_worker_id]
+ for i, part in enumerate(
+ [data[i: i + step] for i in range(0, len(data), step)]
+ )
+ ]
+ processes = []
+ for i in range(n_proc):
+ p = proc(target=_do_parallel_data_prefetch, args=arguments[i])
+ processes += [p]
+
+ # start processes
+ print(f"Start prefetching...")
+ import time
+
+ start = time.time()
+ gather_res = [[] for _ in range(n_proc)]
+ try:
+ for p in processes:
+ p.start()
+
+ k = 0
+ while k < n_proc:
+ # get result
+ res = Q.get()
+ if res == "Done":
+ k += 1
+ else:
+ gather_res[res[0]] = res[1]
+
+ except Exception as e:
+ print("Exception: ", e)
+ for p in processes:
+ p.terminate()
+
+ raise e
+ finally:
+ for p in processes:
+ p.join()
+ print(f"Prefetching complete. [{time.time() - start} sec.]")
+
+ if target_data_type == 'ndarray':
+ if not isinstance(gather_res[0], np.ndarray):
+ return np.concatenate([np.asarray(r) for r in gather_res], axis=0)
+
+ # order outputs
+ return np.concatenate(gather_res, axis=0)
+ elif target_data_type == 'list':
+ out = []
+ for r in gather_res:
+ out.extend(r)
+ return out
+ else:
+ return gather_res
diff --git a/lidm/utils/model_utils.py b/lidm/utils/model_utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..f040e025d44d6c310c874ed7bcef9e789ef73296
--- /dev/null
+++ b/lidm/utils/model_utils.py
@@ -0,0 +1,41 @@
+import os
+
+import torch
+import yaml
+
+from lidm.utils.misc_utils import dict2namespace
+from ..modules.rangenet.model import Model as rangenet
+
+try:
+ from ..modules.spvcnn.model import Model as spvcnn
+ from ..modules.minkowskinet.model import Model as minkowskinet
+except:
+ print('To install torchsparse 1.4.0, please refer to https://github.com/mit-han-lab/torchsparse/tree/74099d10a51c71c14318bce63d6421f698b24f24')
+
+DEFAULT_ROOT = './pretrained_weights'
+
+
+def build_model(dataset_name, model_name, device='cpu'):
+ # config
+ model_folder = os.path.join(DEFAULT_ROOT, dataset_name, model_name)
+
+ if not os.path.isdir(model_folder):
+ raise Exception('Not Available Pretrained Weights!')
+
+ config = yaml.safe_load(open(os.path.join(model_folder, 'config.yaml'), 'r'))
+ if model_name != 'rangenet':
+ config = dict2namespace(config)
+
+ # build model
+ model = eval(model_name)(config)
+
+ # load checkpoint
+ if model_name == 'rangenet':
+ model.load_pretrained_weights(model_folder)
+ else:
+ ckpt = torch.load(os.path.join(model_folder, 'model.ckpt'), map_location="cpu")
+ model.load_state_dict(ckpt['state_dict'], strict=False)
+ model.to(device)
+ model.eval()
+
+ return model
diff --git a/models/first_stage_models/kitti/f_c2_p4/config.yaml b/models/first_stage_models/kitti/f_c2_p4/config.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..fbedeed9f280022b163af049bb2d6f2de8da7e88
--- /dev/null
+++ b/models/first_stage_models/kitti/f_c2_p4/config.yaml
@@ -0,0 +1,57 @@
+model:
+ base_learning_rate: 4.5e-6
+ target: lidm.models.autoencoder.VQModel
+ params:
+ monitor: val/rec_loss
+ embed_dim: 8
+ n_embed: 16384
+ lib_name: lidm
+ use_mask: False # False
+ ddconfig:
+ double_z: false
+ z_channels: 8
+ in_channels: 1
+ out_ch: 1
+ ch: 64
+ ch_mult: [1,2,2,4] # num_down = len(ch_mult)-1
+ strides: [[1,2],[2,2],[2,2]]
+ num_res_blocks: 2
+ attn_levels: []
+ dropout: 0.0
+
+
+data:
+ target: main.DataModuleFromConfig
+ params:
+ batch_size: 4
+ num_workers: 8
+ wrap: true
+ dataset:
+ size: [64, 1024]
+ fov: [ 3,-25 ]
+ depth_range: [ 1.0,56.0 ]
+ depth_scale: 5.84 # np.log2(depth_max + 1)
+ log_scale: true
+ x_range: [ -50.0, 50.0 ]
+ y_range: [ -50.0, 50.0 ]
+ z_range: [ -3.0, 1.0 ]
+ resolution: 1
+ num_channels: 1
+ num_cats: 10
+ num_views: 2
+ num_sem_cats: 19
+ filtered_map_cats: [ ]
+ aug:
+ flip: true
+ rotate: true
+ keypoint_drop: false
+ keypoint_drop_range: [ 5,20 ]
+ randaug: false
+ train:
+ target: lidm.data.kitti.KITTIImageTrain
+ params:
+ condition_key: image
+ validation:
+ target: lidm.data.kitti.KITTIImageValidation
+ params:
+ condition_key: image
diff --git a/models/first_stage_models/kitti/f_c2_p4_wo_logscale/config.yaml b/models/first_stage_models/kitti/f_c2_p4_wo_logscale/config.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..8c4933ad3fe9c1ebba731074487983b1acfe2fb4
--- /dev/null
+++ b/models/first_stage_models/kitti/f_c2_p4_wo_logscale/config.yaml
@@ -0,0 +1,57 @@
+model:
+ base_learning_rate: 4.5e-6
+ target: lidm.models.autoencoder.VQModel
+ params:
+ monitor: val/rec_loss
+ embed_dim: 8
+ n_embed: 16384
+ lib_name: lidm
+ use_mask: False # False
+ ddconfig:
+ double_z: false
+ z_channels: 8
+ in_channels: 1
+ out_ch: 1
+ ch: 64
+ ch_mult: [1,2,2,4] # num_down = len(ch_mult)-1
+ strides: [[1,2],[2,2],[2,2]]
+ num_res_blocks: 2
+ attn_levels: []
+ dropout: 0.0
+
+
+data:
+ target: main.DataModuleFromConfig
+ params:
+ batch_size: 4
+ num_workers: 8
+ wrap: true
+ dataset:
+ size: [64, 1024]
+ fov: [ 3,-25 ]
+ depth_range: [ 1.0,56.0 ]
+ depth_scale: 56 # np.log2(depth_max + 1)
+ log_scale: false
+ x_range: [ -50.0, 50.0 ]
+ y_range: [ -50.0, 50.0 ]
+ z_range: [ -3.0, 1.0 ]
+ resolution: 1
+ num_channels: 1
+ num_cats: 10
+ num_views: 2
+ num_sem_cats: 19
+ filtered_map_cats: [ ]
+ aug:
+ flip: true
+ rotate: true
+ keypoint_drop: false
+ keypoint_drop_range: [ 5,20 ]
+ randaug: false
+ train:
+ target: lidm.data.kitti.KITTIImageTrain
+ params:
+ condition_key: image
+ validation:
+ target: lidm.data.kitti.KITTIImageValidation
+ params:
+ condition_key: image
diff --git a/models/first_stage_models/kitti/f_c2_p4_wo_logscale/model.ckpt b/models/first_stage_models/kitti/f_c2_p4_wo_logscale/model.ckpt
new file mode 100644
index 0000000000000000000000000000000000000000..821d069cbd3698a542eee442ee91977dc47975d5
--- /dev/null
+++ b/models/first_stage_models/kitti/f_c2_p4_wo_logscale/model.ckpt
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:987c413105ad4959899632033091ce9c0a49a56aed1a4f293dcb263ae7022d17
+size 215383923
diff --git a/models/lidm/kitti/cam2lidar/config.yaml b/models/lidm/kitti/cam2lidar/config.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..ef2f705c08fa7f811f96090e9f04ec8298c79fd6
--- /dev/null
+++ b/models/lidm/kitti/cam2lidar/config.yaml
@@ -0,0 +1,110 @@
+model:
+ base_learning_rate: 2.0e-06
+ target: lidm.models.diffusion.ddpm.LatentDiffusion
+ params:
+ linear_start: 0.0015
+ linear_end: 0.0195
+ num_timesteps_cond: 1
+ log_every_t: 100
+ timesteps: 1000
+ image_size: [16, 128]
+ channels: 8
+ monitor: val/loss_simple_ema
+ first_stage_key: image
+ cond_stage_key: camera
+ conditioning_key: crossattn
+ cond_stage_trainable: true
+ verbose: false
+ unet_config:
+ target: lidm.modules.diffusion.openaimodel.UNetModel
+ params:
+ image_size: [16, 128]
+ in_channels: 8
+ out_channels: 8
+ model_channels: 256
+ attention_resolutions: [4, 2, 1]
+ num_res_blocks: 2
+ channel_mult: [1, 2, 4]
+ num_head_channels: 32
+ use_spatial_transformer: true
+ context_dim: 512
+ lib_name: lidm
+ first_stage_config:
+ target: lidm.models.autoencoder.VQModelInterface
+ params:
+ embed_dim: 8
+ n_embed: 16384
+ lib_name: lidm
+ use_mask: False # False
+ ckpt_path: models/first_stage_models/kitti/f_c2_p4_wo_ls/model.ckpt
+ ddconfig:
+ double_z: false
+ z_channels: 8
+ in_channels: 1
+ out_ch: 1
+ ch: 64
+ ch_mult: [1,2,2,4]
+ strides: [[1,2],[2,2],[2,2]]
+ num_res_blocks: 2
+ attn_levels: []
+ dropout: 0.0
+ lossconfig:
+ target: torch.nn.Identity
+ cond_stage_config:
+ target: lidm.modules.encoders.modules.FrozenClipMultiImageEmbedder
+ params:
+ model: ViT-L/14
+ out_dim: 512
+ split_per_view: 4
+
+data:
+ target: main.DataModuleFromConfig
+ params:
+ batch_size: 8
+ num_workers: 8
+ wrap: true
+ dataset:
+ size: [64, 1024]
+ fov: [ 3,-25 ]
+ depth_range: [ 1.0,56.0 ]
+ depth_scale: 56 # np.log2(depth_max + 1)
+ log_scale: false
+ x_range: [ -50.0, 50.0 ]
+ y_range: [ -50.0, 50.0 ]
+ z_range: [ -3.0, 1.0 ]
+ resolution: 1
+ num_channels: 1
+ num_cats: 10
+ num_views: 1
+ num_sem_cats: 19
+ filtered_map_cats: [ ]
+ aug:
+ flip: false
+ rotate: false
+ keypoint_drop: false
+ keypoint_drop_range:
+ randaug: false
+ camera_drop: 0.5
+ train:
+ target: lidm.data.kitti.KITTI360Train
+ params:
+ condition_key: camera
+ split_per_view: 4
+ validation:
+ target: lidm.data.kitti.KITTI360Validation
+ params:
+ condition_key: camera
+ split_per_view: 4
+
+
+lightning:
+ callbacks:
+ image_logger:
+ target: main.ImageLogger
+ params:
+ batch_frequency: 5000
+ max_images: 8
+ increase_log_steps: False
+
+ trainer:
+ benchmark: True
\ No newline at end of file
diff --git a/models/lidm/kitti/cam2lidar/model.ckpt b/models/lidm/kitti/cam2lidar/model.ckpt
new file mode 100644
index 0000000000000000000000000000000000000000..807af1089443089a16864c3fa1275a731576c712
--- /dev/null
+++ b/models/lidm/kitti/cam2lidar/model.ckpt
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:bc6b466d3865ff73813268ba291b32022461e9cfca7cf5bc1534ec08c9cf932a
+size 8093788309
diff --git a/models/lidm/kitti/sem2lidar/config.yaml b/models/lidm/kitti/sem2lidar/config.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..d2cfc8e68c3bbb324570e0e85a2d7181ae78cc43
--- /dev/null
+++ b/models/lidm/kitti/sem2lidar/config.yaml
@@ -0,0 +1,105 @@
+model:
+ base_learning_rate: 1.0e-06
+ target: lidm.models.diffusion.ddpm.LatentDiffusion
+ params:
+ linear_start: 0.0015
+ linear_end: 0.0205
+ num_timesteps_cond: 1
+ log_every_t: 100
+ timesteps: 1000
+ image_size: [16, 128]
+ channels: 8
+ monitor: val/loss_simple_ema
+ first_stage_key: image
+ cond_stage_key: segmentation
+ concat_mode: true
+ cond_stage_trainable: true
+ verbose: false
+ unet_config:
+ target: lidm.modules.diffusion.openaimodel.UNetModel
+ params:
+ image_size: [16, 128]
+ in_channels: 16
+ out_channels: 8
+ model_channels: 256
+ attention_resolutions: [4, 2, 1]
+ num_res_blocks: 2
+ channel_mult: [1, 2, 4]
+ num_head_channels: 32
+ lib_name: lidm
+ first_stage_config:
+ target: lidm.models.autoencoder.VQModelInterface
+ params:
+ embed_dim: 8
+ n_embed: 16384
+ lib_name: lidm
+ use_mask: False # False
+ ckpt_path: models/first_stage_models/kitti/f_c2_p4_wo_ls/model.ckpt
+ ddconfig:
+ double_z: false
+ z_channels: 8
+ in_channels: 1
+ out_ch: 1
+ ch: 64
+ ch_mult: [1,2,2,4]
+ strides: [[1,2],[2,2],[2,2]]
+ num_res_blocks: 2
+ attn_levels: []
+ dropout: 0.0
+ lossconfig:
+ target: torch.nn.Identity
+ cond_stage_config:
+ target: lidm.modules.encoders.modules.SpatialRescaler
+ params:
+ strides: [[1,2],[2,2],[2,2]]
+ in_channels: 20
+ out_channels: 8
+
+data:
+ target: main.DataModuleFromConfig
+ params:
+ batch_size: 16
+ num_workers: 8
+ wrap: true
+ dataset:
+ size: [64, 1024]
+ fov: [ 3,-25 ]
+ depth_range: [ 1.0,56.0 ]
+ depth_scale: 56 # np.log2(depth_max + 1)
+ log_scale: false
+ x_range: [ -50.0, 50.0 ]
+ y_range: [ -50.0, 50.0 ]
+ z_range: [ -3.0, 1.0 ]
+ resolution: 1
+ num_channels: 1
+ num_cats: 10
+ num_views: 2
+ num_sem_cats: 19
+ filtered_map_cats: [ ]
+ aug:
+ flip: true
+ rotate: false
+ keypoint_drop: false
+ keypoint_drop_range: [ 5,20 ]
+ randaug: false
+ train:
+ target: lidm.data.kitti.SemanticKITTITrain
+ params:
+ condition_key: segmentation
+ validation:
+ target: lidm.data.kitti.SemanticKITTIValidation
+ params:
+ condition_key: segmentation
+
+
+lightning:
+ callbacks:
+ image_logger:
+ target: main.ImageLogger
+ params:
+ batch_frequency: 5000
+ max_images: 8
+ increase_log_steps: False
+
+ trainer:
+ benchmark: true
\ No newline at end of file
diff --git a/models/lidm/kitti/text2lidar/config.yaml b/models/lidm/kitti/text2lidar/config.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..e4ceee5ef9a50533fc9b54fd3b9a384109ebeda8
--- /dev/null
+++ b/models/lidm/kitti/text2lidar/config.yaml
@@ -0,0 +1,111 @@
+model:
+ base_learning_rate: 2.0e-06
+ target: lidm.models.diffusion.ddpm.LatentDiffusion
+ params:
+ linear_start: 0.0015
+ linear_end: 0.0195
+ num_timesteps_cond: 1
+ log_every_t: 100
+ timesteps: 1000
+ image_size: [16, 128]
+ channels: 8
+ monitor: val/loss_simple_ema
+ first_stage_key: image
+ cond_stage_key: camera
+ conditioning_key: crossattn
+ cond_stage_trainable: true
+ verbose: false
+ unet_config:
+ target: lidm.modules.diffusion.openaimodel.UNetModel
+ params:
+ image_size: [16, 128]
+ in_channels: 8
+ out_channels: 8
+ model_channels: 256
+ attention_resolutions: [4, 2, 1]
+ num_res_blocks: 2
+ channel_mult: [1, 2, 4]
+ num_head_channels: 32
+ use_spatial_transformer: true
+ context_dim: 512
+ lib_name: lidm
+ first_stage_config:
+ target: lidm.models.autoencoder.VQModelInterface
+ params:
+ embed_dim: 8
+ n_embed: 16384
+ lib_name: lidm
+ use_mask: False # False
+ ckpt_path: models/first_stage_models/kitti/f_c2_p4_wo_ls/model.ckpt
+ ddconfig:
+ double_z: false
+ z_channels: 8
+ in_channels: 1
+ out_ch: 1
+ ch: 64
+ ch_mult: [1,2,2,4]
+ strides: [[1,2],[2,2],[2,2]]
+ num_res_blocks: 2
+ attn_levels: []
+ dropout: 0.0
+ lossconfig:
+ target: torch.nn.Identity
+ cond_stage_config:
+ target: lidm.modules.encoders.modules.FrozenClipMultiImageEmbedder
+ params:
+ model: ViT-L/14
+ split_per_view: 4
+ key: camera
+ out_dim: 512
+
+data:
+ target: main.DataModuleFromConfig
+ params:
+ batch_size: 8
+ num_workers: 8
+ wrap: true
+ dataset:
+ size: [64, 1024]
+ fov: [ 3,-25 ]
+ depth_range: [ 1.0,56.0 ]
+ depth_scale: 56 # np.log2(depth_max + 1)
+ log_scale: false
+ x_range: [ -50.0, 50.0 ]
+ y_range: [ -50.0, 50.0 ]
+ z_range: [ -3.0, 1.0 ]
+ resolution: 1
+ num_channels: 1
+ num_cats: 10
+ num_views: 1
+ num_sem_cats: 19
+ filtered_map_cats: [ ]
+ aug:
+ flip: false
+ rotate: false
+ keypoint_drop: false
+ keypoint_drop_range:
+ randaug: false
+ camera_drop: 0.5
+ train:
+ target: lidm.data.kitti.KITTI360Train
+ params:
+ condition_key: camera
+ split_per_view: 4
+ validation:
+ target: lidm.data.kitti.KITTI360Validation
+ params:
+ condition_key: camera
+ split_per_view: 4
+
+
+lightning:
+ callbacks:
+ image_logger:
+ target: main.ImageLogger
+ params:
+ batch_frequency: 5000
+ max_images: 8
+ increase_log_steps: False
+
+ trainer:
+ benchmark: True
\ No newline at end of file
diff --git a/models/lidm/kitti/uncond/config.yaml b/models/lidm/kitti/uncond/config.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..d2ddd0e4c9916f0a6d4a4335684028e6251a93c2
--- /dev/null
+++ b/models/lidm/kitti/uncond/config.yaml
@@ -0,0 +1,96 @@
+model:
+ base_learning_rate: 1.0e-06
+ target: lidm.models.diffusion.ddpm.LatentDiffusion
+ params:
+ linear_start: 0.0015
+ linear_end: 0.0195
+ num_timesteps_cond: 1
+ log_every_t: 200
+ timesteps: 1000
+ image_size: [16, 128]
+ channels: 8
+ monitor: val/loss_simple_ema
+ first_stage_key: image
+ unet_config:
+ target: lidm.modules.diffusion.openaimodel.UNetModel
+ params:
+ image_size: [16, 128]
+ in_channels: 8
+ out_channels: 8
+ model_channels: 256
+ attention_resolutions: [4, 2, 1]
+ num_res_blocks: 2
+ channel_mult: [1, 2, 4]
+ num_head_channels: 32
+ lib_name: lidm
+ first_stage_config:
+ target: lidm.models.autoencoder.VQModelInterface
+ params:
+ embed_dim: 8
+ n_embed: 16384
+ lib_name: lidm
+ use_mask: False # False
+ ckpt_path: models/first_stage_models/kitti/f_c2_p4/model.ckpt
+ ddconfig:
+ double_z: false
+ z_channels: 8
+ in_channels: 1
+ out_ch: 1
+ ch: 64
+ ch_mult: [1,2,2,4]
+ strides: [[1,2],[2,2],[2,2]]
+ num_res_blocks: 2
+ attn_levels: []
+ dropout: 0.0
+ lossconfig:
+ target: torch.nn.Identity
+ cond_stage_config: "__is_unconditional__"
+
+data:
+ target: main.DataModuleFromConfig
+ params:
+ batch_size: 4
+ num_workers: 8
+ wrap: true
+ dataset:
+ size: [64, 1024]
+ fov: [ 3,-25 ]
+ depth_range: [ 1.0,56.0 ]
+ depth_scale: 5.84 # np.log2(depth_max + 1)
+ log_scale: true
+ x_range: [ -50.0, 50.0 ]
+ y_range: [ -50.0, 50.0 ]
+ z_range: [ -3.0, 1.0 ]
+ resolution: 1
+ num_channels: 1
+ num_cats: 10
+ num_views: 2
+ num_sem_cats: 19
+ filtered_map_cats: [ ]
+ aug:
+ flip: true
+ rotate: false
+ keypoint_drop: false
+ keypoint_drop_range: [ 5,20 ]
+ randaug: false
+ train:
+ target: lidm.data.kitti.KITTI360Train
+ params:
+ condition_key: image
+ validation:
+ target: lidm.data.kitti.KITTI360Validation
+ params:
+ condition_key: image
+
+
+lightning:
+ callbacks:
+ image_logger:
+ target: main.ImageLogger
+ params:
+ batch_frequency: 5000
+ max_images: 8
+ increase_log_steps: False
+
+ trainer:
+ benchmark: true
diff --git a/models/lidm/kitti/uncond_wo_logscale/config.yaml b/models/lidm/kitti/uncond_wo_logscale/config.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..9eb12d7c5c36c8c4dc0ef13f8997df880c39080b
--- /dev/null
+++ b/models/lidm/kitti/uncond_wo_logscale/config.yaml
@@ -0,0 +1,96 @@
+model:
+ base_learning_rate: 1.0e-06
+ target: lidm.models.diffusion.ddpm.LatentDiffusion
+ params:
+ linear_start: 0.0015
+ linear_end: 0.0195
+ num_timesteps_cond: 1
+ log_every_t: 200
+ timesteps: 1000
+ image_size: [16, 128]
+ channels: 8
+ monitor: val/loss_simple_ema
+ first_stage_key: image
+ unet_config:
+ target: lidm.modules.diffusion.openaimodel.UNetModel
+ params:
+ image_size: [16, 128]
+ in_channels: 8
+ out_channels: 8
+ model_channels: 256
+ attention_resolutions: [4, 2, 1]
+ num_res_blocks: 2
+ channel_mult: [1, 2, 4]
+ num_head_channels: 32
+ lib_name: lidm
+ first_stage_config:
+ target: lidm.models.autoencoder.VQModelInterface
+ params:
+ embed_dim: 8
+ n_embed: 16384
+ lib_name: lidm
+ use_mask: False # False
+ ckpt_path: models/first_stage_models/kitti/f_c2_p4_wo_ls/model.ckpt
+ ddconfig:
+ double_z: false
+ z_channels: 8
+ in_channels: 1
+ out_ch: 1
+ ch: 64
+ ch_mult: [1,2,2,4]
+ strides: [[1,2],[2,2],[2,2]]
+ num_res_blocks: 2
+ attn_levels: []
+ dropout: 0.0
+ lossconfig:
+ target: torch.nn.Identity
+ cond_stage_config: "__is_unconditional__"
+
+data:
+ target: main.DataModuleFromConfig
+ params:
+ batch_size: 4
+ num_workers: 8
+ wrap: true
+ dataset:
+ size: [64, 1024]
+ fov: [ 3,-25 ]
+ depth_range: [ 1.0,56.0 ]
+ depth_scale: 56 # np.log2(depth_max + 1)
+ log_scale: false
+ x_range: [ -50.0, 50.0 ]
+ y_range: [ -50.0, 50.0 ]
+ z_range: [ -3.0, 1.0 ]
+ resolution: 1
+ num_channels: 1
+ num_cats: 10
+ num_views: 2
+ num_sem_cats: 19
+ filtered_map_cats: [ ]
+ aug:
+ flip: true
+ rotate: false
+ keypoint_drop: false
+ keypoint_drop_range: [ 5,20 ]
+ randaug: false
+ train:
+ target: lidm.data.kitti.KITTI360Train
+ params:
+ condition_key: image
+ validation:
+ target: lidm.data.kitti.KITTI360Validation
+ params:
+ condition_key: image
+
+
+lightning:
+ callbacks:
+ image_logger:
+ target: main.ImageLogger
+ params:
+ batch_frequency: 5000
+ max_images: 8
+ increase_log_steps: False
+
+ trainer:
+ benchmark: true
diff --git a/sample_cond.py b/sample_cond.py
new file mode 100644
index 0000000000000000000000000000000000000000..4609b4cd88d7ad8a90000b0e901d457352c88f54
--- /dev/null
+++ b/sample_cond.py
@@ -0,0 +1,109 @@
+import os
+import torch
+import numpy as np
+
+from omegaconf import OmegaConf
+from PIL import Image
+
+from lidm.models.diffusion.ddim import DDIMSampler
+from lidm.utils.misc_utils import instantiate_from_config, isimage, ismap
+from lidm.utils.lidar_utils import range2pcd
+from app_config import DEVICE
+
+
+CUSTOM_STEPS = 50
+ETA = 1.0
+
+# model loading
+MODEL_PATH = './models/lidm/kitti/cam2lidar'
+CFG_PATH = os.path.join(MODEL_PATH, 'config.yaml')
+CKPT_PATH = os.path.join(MODEL_PATH, 'model.ckpt')
+
+# settings
+model_config = OmegaConf.load(CFG_PATH)
+
+
+def custom_to_pcd(x, config, rgb=None):
+ x = x.squeeze().detach().cpu().numpy()
+ x = (np.clip(x, -1., 1.) + 1.) / 2.
+ if rgb is not None:
+ rgb = rgb.squeeze().detach().cpu().numpy()
+ rgb = (np.clip(rgb, -1., 1.) + 1.) / 2.
+ rgb = rgb.transpose(1, 2, 0)
+ xyz, rgb, _ = range2pcd(x, color=rgb, **config['data']['params']['dataset'])
+
+ return xyz, rgb
+
+
+def custom_to_pil(x):
+ x = x.detach().cpu().squeeze().numpy()
+ x = (np.clip(x, -1., 1.) + 1.) / 2.
+ x = (255 * x).astype(np.uint8)
+
+ if x.ndim == 3:
+ x = x.transpose(1, 2, 0)
+ x = Image.fromarray(x)
+
+ return x
+
+
+def logs2pil(logs, keys=["sample"]):
+ imgs = dict()
+ for k in logs:
+ try:
+ if len(logs[k].shape) == 4:
+ img = custom_to_pil(logs[k][0, ...])
+ elif len(logs[k].shape) == 3:
+ img = custom_to_pil(logs[k])
+ else:
+ print(f"Unknown format for key {k}. ")
+ img = None
+ except:
+ img = None
+ imgs[k] = img
+ return imgs
+
+
+def load_model_from_config(config, sd, device):
+ model = instantiate_from_config(config)
+ model.load_state_dict(sd, strict=False)
+ model.to(device)
+ model.eval()
+ return model
+
+
+def load_model():
+ pl_sd = torch.load(CKPT_PATH, map_location="cpu")
+ model = load_model_from_config(model_config.model, pl_sd["state_dict"], DEVICE)
+ return model
+
+
+@torch.no_grad()
+def convsample_ddim(model, cond, steps, shape, eta=1.0, verbose=False):
+ ddim = DDIMSampler(model)
+ bs = shape[0]
+ shape = shape[1:]
+ samples, intermediates = ddim.sample(steps, conditioning=cond, batch_size=bs, shape=shape, eta=eta, verbose=verbose, disable_tqdm=True)
+ return samples, intermediates
+
+
+@torch.no_grad()
+def make_convolutional_sample(model, batch, batch_size, custom_steps=None, eta=1.0):
+ xc = batch['camera']
+ c = model.get_learned_conditioning(xc.to(model.device))
+
+ with model.ema_scope("Plotting"):
+ samples, z_denoise_row = model.sample_log(cond=c, batch_size=batch_size, ddim=True,
+ ddim_steps=custom_steps, eta=eta)
+ x_samples = model.decode_first_stage(samples)
+
+ return x_samples
+
+
+def sample(model, cond):
+ batch = {'camera': cond}
+ img = make_convolutional_sample(model, batch, batch_size=1, custom_steps=CUSTOM_STEPS, eta=ETA) # TODO add arguments for batch_size, custom_steps and eta
+ img = img[0, 0]
+ pcd = custom_to_pcd(img, model_config)[0].astype(np.float32)
+ return img, pcd
+