diff --git a/.idea/.name b/.idea/.name deleted file mode 100644 index 42061c0..0000000 --- a/.idea/.name +++ /dev/null @@ -1 +0,0 @@ -README.md \ No newline at end of file diff --git a/.idea/CCSR.iml b/.idea/CCSR.iml deleted file mode 100644 index d0876a7..0000000 --- a/.idea/CCSR.iml +++ /dev/null @@ -1,8 +0,0 @@ - - - - - - - - \ No newline at end of file diff --git a/.idea/inspectionProfiles/Project_Default.xml b/.idea/inspectionProfiles/Project_Default.xml deleted file mode 100644 index 0c6e339..0000000 --- a/.idea/inspectionProfiles/Project_Default.xml +++ /dev/null @@ -1,14 +0,0 @@ - - - - \ No newline at end of file diff --git a/.idea/inspectionProfiles/profiles_settings.xml b/.idea/inspectionProfiles/profiles_settings.xml deleted file mode 100644 index 105ce2d..0000000 --- a/.idea/inspectionProfiles/profiles_settings.xml +++ /dev/null @@ -1,6 +0,0 @@ - - - - \ No newline at end of file diff --git a/.idea/misc.xml b/.idea/misc.xml deleted file mode 100644 index d56657a..0000000 --- a/.idea/misc.xml +++ /dev/null @@ -1,4 +0,0 @@ - - - - \ No newline at end of file diff --git a/.idea/modules.xml b/.idea/modules.xml deleted file mode 100644 index bdebacb..0000000 --- a/.idea/modules.xml +++ /dev/null @@ -1,8 +0,0 @@ - - - - - - - - \ No newline at end of file diff --git a/.idea/vcs.xml b/.idea/vcs.xml deleted file mode 100644 index f6eb092..0000000 --- a/.idea/vcs.xml +++ /dev/null @@ -1,7 +0,0 @@ - - - - - - - \ No newline at end of file diff --git a/.idea/workspace.xml b/.idea/workspace.xml deleted file mode 100644 index e17ae08..0000000 --- a/.idea/workspace.xml +++ /dev/null @@ -1,39 +0,0 @@ - - - - - - - - - - - - - - { - "keyToString": { - "RunOnceActivity.OpenProjectViewOnStart": "true", - "RunOnceActivity.ShowReadmeOnStart": "true" - } -} - - - - - 1703257984114 - - - - \ No newline at end of file diff --git a/model/q_sampler.py b/model/q_sampler.py index 8a78c5b..bf5e056 100644 --- a/model/q_sampler.py +++ b/model/q_sampler.py @@ -3,13 +3,10 @@ from typing import Optional, Tuple, Dict, List, Callable import torch import numpy as np from tqdm import tqdm -import einops -import os -from PIL import Image from ..ldm.modules.diffusionmodules.util import make_beta_schedule from ..model.cond_fn import Guidance -from ..utils.image import ( +from ..utils.align_color import ( wavelet_reconstruction, adaptive_instance_normalization ) diff --git a/model/spaced_sampler.py b/model/spaced_sampler.py index fee3bff..f82fee7 100644 --- a/model/spaced_sampler.py +++ b/model/spaced_sampler.py @@ -9,7 +9,7 @@ from PIL import Image from ..ldm.modules.diffusionmodules.util import make_beta_schedule from ..model.cond_fn import Guidance -from ..utils.image import ( +from ..utils.align_color import ( wavelet_reconstruction, adaptive_instance_normalization ) diff --git a/requirements.txt b/requirements.txt index 273dad9..59b5d27 100644 --- a/requirements.txt +++ b/requirements.txt @@ -2,4 +2,5 @@ taming-transformers omegaconf einops pytorch-lightning -lightning-fabric \ No newline at end of file +lightning-fabric +open-clip-torch \ No newline at end of file diff --git a/taming/data/ade20k.py b/taming/data/ade20k.py deleted file mode 100644 index 366dae9..0000000 --- a/taming/data/ade20k.py +++ /dev/null @@ -1,124 +0,0 @@ -import os -import numpy as np -import cv2 -import albumentations -from PIL import Image -from torch.utils.data import Dataset - -from taming.data.sflckr import SegmentationBase # for examples included in repo - - -class Examples(SegmentationBase): - def __init__(self, size=256, random_crop=False, interpolation="bicubic"): - super().__init__(data_csv="data/ade20k_examples.txt", - data_root="data/ade20k_images", - segmentation_root="data/ade20k_segmentations", - size=size, random_crop=random_crop, - interpolation=interpolation, - n_labels=151, shift_segmentation=False) - - -# With semantic map and scene label -class ADE20kBase(Dataset): - def __init__(self, config=None, size=None, random_crop=False, interpolation="bicubic", crop_size=None): - self.split = self.get_split() - self.n_labels = 151 # unknown + 150 - self.data_csv = {"train": "data/ade20k_train.txt", - "validation": "data/ade20k_test.txt"}[self.split] - self.data_root = "data/ade20k_root" - with open(os.path.join(self.data_root, "sceneCategories.txt"), "r") as f: - self.scene_categories = f.read().splitlines() - self.scene_categories = dict(line.split() for line in self.scene_categories) - with open(self.data_csv, "r") as f: - self.image_paths = f.read().splitlines() - self._length = len(self.image_paths) - self.labels = { - "relative_file_path_": [l for l in self.image_paths], - "file_path_": [os.path.join(self.data_root, "images", l) - for l in self.image_paths], - "relative_segmentation_path_": [l.replace(".jpg", ".png") - for l in self.image_paths], - "segmentation_path_": [os.path.join(self.data_root, "annotations", - l.replace(".jpg", ".png")) - for l in self.image_paths], - "scene_category": [self.scene_categories[l.split("/")[1].replace(".jpg", "")] - for l in self.image_paths], - } - - size = None if size is not None and size<=0 else size - self.size = size - if crop_size is None: - self.crop_size = size if size is not None else None - else: - self.crop_size = crop_size - if self.size is not None: - self.interpolation = interpolation - self.interpolation = { - "nearest": cv2.INTER_NEAREST, - "bilinear": cv2.INTER_LINEAR, - "bicubic": cv2.INTER_CUBIC, - "area": cv2.INTER_AREA, - "lanczos": cv2.INTER_LANCZOS4}[self.interpolation] - self.image_rescaler = albumentations.SmallestMaxSize(max_size=self.size, - interpolation=self.interpolation) - self.segmentation_rescaler = albumentations.SmallestMaxSize(max_size=self.size, - interpolation=cv2.INTER_NEAREST) - - if crop_size is not None: - self.center_crop = not random_crop - if self.center_crop: - self.cropper = albumentations.CenterCrop(height=self.crop_size, width=self.crop_size) - else: - self.cropper = albumentations.RandomCrop(height=self.crop_size, width=self.crop_size) - self.preprocessor = self.cropper - - def __len__(self): - return self._length - - def __getitem__(self, i): - example = dict((k, self.labels[k][i]) for k in self.labels) - image = Image.open(example["file_path_"]) - if not image.mode == "RGB": - image = image.convert("RGB") - image = np.array(image).astype(np.uint8) - if self.size is not None: - image = self.image_rescaler(image=image)["image"] - segmentation = Image.open(example["segmentation_path_"]) - segmentation = np.array(segmentation).astype(np.uint8) - if self.size is not None: - segmentation = self.segmentation_rescaler(image=segmentation)["image"] - if self.size is not None: - processed = self.preprocessor(image=image, mask=segmentation) - else: - processed = {"image": image, "mask": segmentation} - example["image"] = (processed["image"]/127.5 - 1.0).astype(np.float32) - segmentation = processed["mask"] - onehot = np.eye(self.n_labels)[segmentation] - example["segmentation"] = onehot - return example - - -class ADE20kTrain(ADE20kBase): - # default to random_crop=True - def __init__(self, config=None, size=None, random_crop=True, interpolation="bicubic", crop_size=None): - super().__init__(config=config, size=size, random_crop=random_crop, - interpolation=interpolation, crop_size=crop_size) - - def get_split(self): - return "train" - - -class ADE20kValidation(ADE20kBase): - def get_split(self): - return "validation" - - -if __name__ == "__main__": - dset = ADE20kValidation() - ex = dset[0] - for k in ["image", "scene_category", "segmentation"]: - print(type(ex[k])) - try: - print(ex[k].shape) - except: - print(ex[k]) diff --git a/taming/data/annotated_objects_coco.py b/taming/data/annotated_objects_coco.py deleted file mode 100644 index af000ec..0000000 --- a/taming/data/annotated_objects_coco.py +++ /dev/null @@ -1,139 +0,0 @@ -import json -from itertools import chain -from pathlib import Path -from typing import Iterable, Dict, List, Callable, Any -from collections import defaultdict - -from tqdm import tqdm - -from taming.data.annotated_objects_dataset import AnnotatedObjectsDataset -from taming.data.helper_types import Annotation, ImageDescription, Category - -COCO_PATH_STRUCTURE = { - 'train': { - 'top_level': '', - 'instances_annotations': 'annotations/instances_train2017.json', - 'stuff_annotations': 'annotations/stuff_train2017.json', - 'files': 'train2017' - }, - 'validation': { - 'top_level': '', - 'instances_annotations': 'annotations/instances_val2017.json', - 'stuff_annotations': 'annotations/stuff_val2017.json', - 'files': 'val2017' - } -} - - -def load_image_descriptions(description_json: List[Dict]) -> Dict[str, ImageDescription]: - return { - str(img['id']): ImageDescription( - id=img['id'], - license=img.get('license'), - file_name=img['file_name'], - coco_url=img['coco_url'], - original_size=(img['width'], img['height']), - date_captured=img.get('date_captured'), - flickr_url=img.get('flickr_url') - ) - for img in description_json - } - - -def load_categories(category_json: Iterable) -> Dict[str, Category]: - return {str(cat['id']): Category(id=str(cat['id']), super_category=cat['supercategory'], name=cat['name']) - for cat in category_json if cat['name'] != 'other'} - - -def load_annotations(annotations_json: List[Dict], image_descriptions: Dict[str, ImageDescription], - category_no_for_id: Callable[[str], int], split: str) -> Dict[str, List[Annotation]]: - annotations = defaultdict(list) - total = sum(len(a) for a in annotations_json) - for ann in tqdm(chain(*annotations_json), f'Loading {split} annotations', total=total): - image_id = str(ann['image_id']) - if image_id not in image_descriptions: - raise ValueError(f'image_id [{image_id}] has no image description.') - category_id = ann['category_id'] - try: - category_no = category_no_for_id(str(category_id)) - except KeyError: - continue - - width, height = image_descriptions[image_id].original_size - bbox = (ann['bbox'][0] / width, ann['bbox'][1] / height, ann['bbox'][2] / width, ann['bbox'][3] / height) - - annotations[image_id].append( - Annotation( - id=ann['id'], - area=bbox[2]*bbox[3], # use bbox area - is_group_of=ann['iscrowd'], - image_id=ann['image_id'], - bbox=bbox, - category_id=str(category_id), - category_no=category_no - ) - ) - return dict(annotations) - - -class AnnotatedObjectsCoco(AnnotatedObjectsDataset): - def __init__(self, use_things: bool = True, use_stuff: bool = True, **kwargs): - """ - @param data_path: is the path to the following folder structure: - coco/ - ├── annotations - │ ├── instances_train2017.json - │ ├── instances_val2017.json - │ ├── stuff_train2017.json - │ └── stuff_val2017.json - ├── train2017 - │ ├── 000000000009.jpg - │ ├── 000000000025.jpg - │ └── ... - ├── val2017 - │ ├── 000000000139.jpg - │ ├── 000000000285.jpg - │ └── ... - @param: split: one of 'train' or 'validation' - @param: desired image size (give square images) - """ - super().__init__(**kwargs) - self.use_things = use_things - self.use_stuff = use_stuff - - with open(self.paths['instances_annotations']) as f: - inst_data_json = json.load(f) - with open(self.paths['stuff_annotations']) as f: - stuff_data_json = json.load(f) - - category_jsons = [] - annotation_jsons = [] - if self.use_things: - category_jsons.append(inst_data_json['categories']) - annotation_jsons.append(inst_data_json['annotations']) - if self.use_stuff: - category_jsons.append(stuff_data_json['categories']) - annotation_jsons.append(stuff_data_json['annotations']) - - self.categories = load_categories(chain(*category_jsons)) - self.filter_categories() - self.setup_category_id_and_number() - - self.image_descriptions = load_image_descriptions(inst_data_json['images']) - annotations = load_annotations(annotation_jsons, self.image_descriptions, self.get_category_number, self.split) - self.annotations = self.filter_object_number(annotations, self.min_object_area, - self.min_objects_per_image, self.max_objects_per_image) - self.image_ids = list(self.annotations.keys()) - self.clean_up_annotations_and_image_descriptions() - - def get_path_structure(self) -> Dict[str, str]: - if self.split not in COCO_PATH_STRUCTURE: - raise ValueError(f'Split [{self.split} does not exist for COCO data.]') - return COCO_PATH_STRUCTURE[self.split] - - def get_image_path(self, image_id: str) -> Path: - return self.paths['files'].joinpath(self.image_descriptions[str(image_id)].file_name) - - def get_image_description(self, image_id: str) -> Dict[str, Any]: - # noinspection PyProtectedMember - return self.image_descriptions[image_id]._asdict() diff --git a/taming/data/annotated_objects_dataset.py b/taming/data/annotated_objects_dataset.py deleted file mode 100644 index 53cc346..0000000 --- a/taming/data/annotated_objects_dataset.py +++ /dev/null @@ -1,218 +0,0 @@ -from pathlib import Path -from typing import Optional, List, Callable, Dict, Any, Union -import warnings - -import PIL.Image as pil_image -from torch import Tensor -from torch.utils.data import Dataset -from torchvision import transforms - -from taming.data.conditional_builder.objects_bbox import ObjectsBoundingBoxConditionalBuilder -from taming.data.conditional_builder.objects_center_points import ObjectsCenterPointsConditionalBuilder -from taming.data.conditional_builder.utils import load_object_from_string -from taming.data.helper_types import BoundingBox, CropMethodType, Image, Annotation, SplitType -from taming.data.image_transforms import CenterCropReturnCoordinates, RandomCrop1dReturnCoordinates, \ - Random2dCropReturnCoordinates, RandomHorizontalFlipReturn, convert_pil_to_tensor - - -class AnnotatedObjectsDataset(Dataset): - def __init__(self, data_path: Union[str, Path], split: SplitType, keys: List[str], target_image_size: int, - min_object_area: float, min_objects_per_image: int, max_objects_per_image: int, - crop_method: CropMethodType, random_flip: bool, no_tokens: int, use_group_parameter: bool, - encode_crop: bool, category_allow_list_target: str = "", category_mapping_target: str = "", - no_object_classes: Optional[int] = None): - self.data_path = data_path - self.split = split - self.keys = keys - self.target_image_size = target_image_size - self.min_object_area = min_object_area - self.min_objects_per_image = min_objects_per_image - self.max_objects_per_image = max_objects_per_image - self.crop_method = crop_method - self.random_flip = random_flip - self.no_tokens = no_tokens - self.use_group_parameter = use_group_parameter - self.encode_crop = encode_crop - - self.annotations = None - self.image_descriptions = None - self.categories = None - self.category_ids = None - self.category_number = None - self.image_ids = None - self.transform_functions: List[Callable] = self.setup_transform(target_image_size, crop_method, random_flip) - self.paths = self.build_paths(self.data_path) - self._conditional_builders = None - self.category_allow_list = None - if category_allow_list_target: - allow_list = load_object_from_string(category_allow_list_target) - self.category_allow_list = {name for name, _ in allow_list} - self.category_mapping = {} - if category_mapping_target: - self.category_mapping = load_object_from_string(category_mapping_target) - self.no_object_classes = no_object_classes - - def build_paths(self, top_level: Union[str, Path]) -> Dict[str, Path]: - top_level = Path(top_level) - sub_paths = {name: top_level.joinpath(sub_path) for name, sub_path in self.get_path_structure().items()} - for path in sub_paths.values(): - if not path.exists(): - raise FileNotFoundError(f'{type(self).__name__} data structure error: [{path}] does not exist.') - return sub_paths - - @staticmethod - def load_image_from_disk(path: Path) -> Image: - return pil_image.open(path).convert('RGB') - - @staticmethod - def setup_transform(target_image_size: int, crop_method: CropMethodType, random_flip: bool): - transform_functions = [] - if crop_method == 'none': - transform_functions.append(transforms.Resize((target_image_size, target_image_size))) - elif crop_method == 'center': - transform_functions.extend([ - transforms.Resize(target_image_size), - CenterCropReturnCoordinates(target_image_size) - ]) - elif crop_method == 'random-1d': - transform_functions.extend([ - transforms.Resize(target_image_size), - RandomCrop1dReturnCoordinates(target_image_size) - ]) - elif crop_method == 'random-2d': - transform_functions.extend([ - Random2dCropReturnCoordinates(target_image_size), - transforms.Resize(target_image_size) - ]) - elif crop_method is None: - return None - else: - raise ValueError(f'Received invalid crop method [{crop_method}].') - if random_flip: - transform_functions.append(RandomHorizontalFlipReturn()) - transform_functions.append(transforms.Lambda(lambda x: x / 127.5 - 1.)) - return transform_functions - - def image_transform(self, x: Tensor) -> (Optional[BoundingBox], Optional[bool], Tensor): - crop_bbox = None - flipped = None - for t in self.transform_functions: - if isinstance(t, (RandomCrop1dReturnCoordinates, CenterCropReturnCoordinates, Random2dCropReturnCoordinates)): - crop_bbox, x = t(x) - elif isinstance(t, RandomHorizontalFlipReturn): - flipped, x = t(x) - else: - x = t(x) - return crop_bbox, flipped, x - - @property - def no_classes(self) -> int: - return self.no_object_classes if self.no_object_classes else 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 = { - 'objects_center_points': ObjectsCenterPointsConditionalBuilder( - self.no_classes, - self.max_objects_per_image, - self.no_tokens, - self.encode_crop, - self.use_group_parameter, - getattr(self, 'use_additional_parameters', False) - ), - 'objects_bbox': ObjectsBoundingBoxConditionalBuilder( - self.no_classes, - self.max_objects_per_image, - self.no_tokens, - self.encode_crop, - self.use_group_parameter, - getattr(self, 'use_additional_parameters', False) - ) - } - return self._conditional_builders - - def filter_categories(self) -> None: - if self.category_allow_list: - self.categories = {id_: cat for id_, cat in self.categories.items() if cat.name in self.category_allow_list} - if self.category_mapping: - self.categories = {id_: cat for id_, cat in self.categories.items() if cat.id not in self.category_mapping} - - def setup_category_id_and_number(self) -> None: - self.category_ids = list(self.categories.keys()) - self.category_ids.sort() - if '/m/01s55n' in self.category_ids: - self.category_ids.remove('/m/01s55n') - self.category_ids.append('/m/01s55n') - self.category_number = {category_id: i for i, category_id in enumerate(self.category_ids)} - if self.category_allow_list is not None and self.category_mapping is None \ - and len(self.category_ids) != len(self.category_allow_list): - warnings.warn('Unexpected number of categories: Mismatch with category_allow_list. ' - 'Make sure all names in category_allow_list exist.') - - def clean_up_annotations_and_image_descriptions(self) -> None: - image_id_set = set(self.image_ids) - self.annotations = {k: v for k, v in self.annotations.items() if k in image_id_set} - self.image_descriptions = {k: v for k, v in self.image_descriptions.items() if k in image_id_set} - - @staticmethod - def filter_object_number(all_annotations: Dict[str, List[Annotation]], min_object_area: float, - min_objects_per_image: int, max_objects_per_image: int) -> Dict[str, List[Annotation]]: - filtered = {} - for image_id, annotations in all_annotations.items(): - annotations_with_min_area = [a for a in annotations if a.area > min_object_area] - if min_objects_per_image <= len(annotations_with_min_area) <= max_objects_per_image: - filtered[image_id] = annotations_with_min_area - return filtered - - def __len__(self): - return len(self.image_ids) - - def __getitem__(self, n: int) -> Dict[str, Any]: - image_id = self.get_image_id(n) - sample = self.get_image_description(image_id) - sample['annotations'] = self.get_annotation(image_id) - - if 'image' in self.keys: - sample['image_path'] = str(self.get_image_path(image_id)) - sample['image'] = self.load_image_from_disk(sample['image_path']) - sample['image'] = convert_pil_to_tensor(sample['image']) - sample['crop_bbox'], sample['flipped'], sample['image'] = self.image_transform(sample['image']) - sample['image'] = sample['image'].permute(1, 2, 0) - - for conditional, builder in self.conditional_builders.items(): - if conditional in self.keys: - sample[conditional] = builder.build(sample['annotations'], sample['crop_bbox'], sample['flipped']) - - if self.keys: - # only return specified keys - sample = {key: sample[key] for key in self.keys} - return sample - - def get_image_id(self, no: int) -> str: - return self.image_ids[no] - - def get_annotation(self, image_id: str) -> str: - return self.annotations[image_id] - - def get_textual_label_for_category_id(self, category_id: str) -> str: - return self.categories[category_id].name - - def get_textual_label_for_category_no(self, category_no: int) -> str: - return self.categories[self.get_category_id(category_no)].name - - def get_category_number(self, category_id: str) -> int: - return self.category_number[category_id] - - def get_category_id(self, category_no: int) -> str: - return self.category_ids[category_no] - - def get_image_description(self, image_id: str) -> Dict[str, Any]: - raise NotImplementedError() - - def get_path_structure(self): - raise NotImplementedError - - def get_image_path(self, image_id: str) -> Path: - raise NotImplementedError diff --git a/taming/data/annotated_objects_open_images.py b/taming/data/annotated_objects_open_images.py deleted file mode 100644 index aede680..0000000 --- a/taming/data/annotated_objects_open_images.py +++ /dev/null @@ -1,137 +0,0 @@ -from collections import defaultdict -from csv import DictReader, reader as TupleReader -from pathlib import Path -from typing import Dict, List, Any -import warnings - -from taming.data.annotated_objects_dataset import AnnotatedObjectsDataset -from taming.data.helper_types import Annotation, Category -from tqdm import tqdm - -OPEN_IMAGES_STRUCTURE = { - 'train': { - 'top_level': '', - 'class_descriptions': 'class-descriptions-boxable.csv', - 'annotations': 'oidv6-train-annotations-bbox.csv', - 'file_list': 'train-images-boxable.csv', - 'files': 'train' - }, - 'validation': { - 'top_level': '', - 'class_descriptions': 'class-descriptions-boxable.csv', - 'annotations': 'validation-annotations-bbox.csv', - 'file_list': 'validation-images.csv', - 'files': 'validation' - }, - 'test': { - 'top_level': '', - 'class_descriptions': 'class-descriptions-boxable.csv', - 'annotations': 'test-annotations-bbox.csv', - 'file_list': 'test-images.csv', - 'files': 'test' - } -} - - -def load_annotations(descriptor_path: Path, min_object_area: float, category_mapping: Dict[str, str], - category_no_for_id: Dict[str, int]) -> Dict[str, List[Annotation]]: - annotations: Dict[str, List[Annotation]] = defaultdict(list) - with open(descriptor_path) as file: - reader = DictReader(file) - for i, row in tqdm(enumerate(reader), total=14620000, desc='Loading OpenImages annotations'): - width = float(row['XMax']) - float(row['XMin']) - height = float(row['YMax']) - float(row['YMin']) - area = width * height - category_id = row['LabelName'] - if category_id in category_mapping: - category_id = category_mapping[category_id] - if area >= min_object_area and category_id in category_no_for_id: - annotations[row['ImageID']].append( - Annotation( - id=i, - image_id=row['ImageID'], - source=row['Source'], - category_id=category_id, - category_no=category_no_for_id[category_id], - confidence=float(row['Confidence']), - bbox=(float(row['XMin']), float(row['YMin']), width, height), - area=area, - is_occluded=bool(int(row['IsOccluded'])), - is_truncated=bool(int(row['IsTruncated'])), - is_group_of=bool(int(row['IsGroupOf'])), - is_depiction=bool(int(row['IsDepiction'])), - is_inside=bool(int(row['IsInside'])) - ) - ) - if 'train' in str(descriptor_path) and i < 14000000: - warnings.warn(f'Running with subset of Open Images. Train dataset has length [{len(annotations)}].') - return dict(annotations) - - -def load_image_ids(csv_path: Path) -> List[str]: - with open(csv_path) as file: - reader = DictReader(file) - return [row['image_name'] for row in reader] - - -def load_categories(csv_path: Path) -> Dict[str, Category]: - with open(csv_path) as file: - reader = TupleReader(file) - return {row[0]: Category(id=row[0], name=row[1], super_category=None) for row in reader} - - -class AnnotatedObjectsOpenImages(AnnotatedObjectsDataset): - def __init__(self, use_additional_parameters: bool, **kwargs): - """ - @param data_path: is the path to the following folder structure: - open_images/ - │ oidv6-train-annotations-bbox.csv - ├── class-descriptions-boxable.csv - ├── oidv6-train-annotations-bbox.csv - ├── test - │ ├── 000026e7ee790996.jpg - │ ├── 000062a39995e348.jpg - │ └── ... - ├── test-annotations-bbox.csv - ├── test-images.csv - ├── train - │ ├── 000002b66c9c498e.jpg - │ ├── 000002b97e5471a0.jpg - │ └── ... - ├── train-images-boxable.csv - ├── validation - │ ├── 0001eeaf4aed83f9.jpg - │ ├── 0004886b7d043cfd.jpg - │ └── ... - ├── validation-annotations-bbox.csv - └── validation-images.csv - @param: split: one of 'train', 'validation' or 'test' - @param: desired image size (returns square images) - """ - - super().__init__(**kwargs) - self.use_additional_parameters = use_additional_parameters - - self.categories = load_categories(self.paths['class_descriptions']) - self.filter_categories() - self.setup_category_id_and_number() - - self.image_descriptions = {} - annotations = load_annotations(self.paths['annotations'], self.min_object_area, self.category_mapping, - self.category_number) - self.annotations = self.filter_object_number(annotations, self.min_object_area, self.min_objects_per_image, - self.max_objects_per_image) - self.image_ids = list(self.annotations.keys()) - self.clean_up_annotations_and_image_descriptions() - - def get_path_structure(self) -> Dict[str, str]: - if self.split not in OPEN_IMAGES_STRUCTURE: - raise ValueError(f'Split [{self.split} does not exist for Open Images data.]') - return OPEN_IMAGES_STRUCTURE[self.split] - - def get_image_path(self, image_id: str) -> Path: - return self.paths['files'].joinpath(f'{image_id:0>16}.jpg') - - def get_image_description(self, image_id: str) -> Dict[str, Any]: - image_path = self.get_image_path(image_id) - return {'file_path': str(image_path), 'file_name': image_path.name} diff --git a/taming/data/base.py b/taming/data/base.py deleted file mode 100644 index e21667d..0000000 --- a/taming/data/base.py +++ /dev/null @@ -1,70 +0,0 @@ -import bisect -import numpy as np -import albumentations -from PIL import Image -from torch.utils.data import Dataset, ConcatDataset - - -class ConcatDatasetWithIndex(ConcatDataset): - """Modified from original pytorch code to return dataset idx""" - def __getitem__(self, idx): - if idx < 0: - if -idx > len(self): - raise ValueError("absolute value of index should not exceed dataset length") - idx = len(self) + idx - dataset_idx = bisect.bisect_right(self.cumulative_sizes, idx) - if dataset_idx == 0: - sample_idx = idx - else: - sample_idx = idx - self.cumulative_sizes[dataset_idx - 1] - return self.datasets[dataset_idx][sample_idx], dataset_idx - - -class ImagePaths(Dataset): - def __init__(self, paths, size=None, random_crop=False, labels=None): - self.size = size - self.random_crop = random_crop - - self.labels = dict() if labels is None else labels - self.labels["file_path_"] = paths - self._length = len(paths) - - if self.size is not None and self.size > 0: - self.rescaler = albumentations.SmallestMaxSize(max_size = self.size) - if not self.random_crop: - self.cropper = albumentations.CenterCrop(height=self.size,width=self.size) - else: - self.cropper = albumentations.RandomCrop(height=self.size,width=self.size) - self.preprocessor = albumentations.Compose([self.rescaler, self.cropper]) - else: - self.preprocessor = lambda **kwargs: kwargs - - def __len__(self): - return self._length - - def preprocess_image(self, image_path): - image = Image.open(image_path) - if not image.mode == "RGB": - image = image.convert("RGB") - image = np.array(image).astype(np.uint8) - image = self.preprocessor(image=image)["image"] - image = (image/127.5 - 1.0).astype(np.float32) - return image - - def __getitem__(self, i): - example = dict() - example["image"] = self.preprocess_image(self.labels["file_path_"][i]) - for k in self.labels: - example[k] = self.labels[k][i] - return example - - -class NumpyPaths(ImagePaths): - def preprocess_image(self, image_path): - image = np.load(image_path).squeeze(0) # 3 x 1024 x 1024 - image = np.transpose(image, (1,2,0)) - image = Image.fromarray(image, mode="RGB") - image = np.array(image).astype(np.uint8) - image = self.preprocessor(image=image)["image"] - image = (image/127.5 - 1.0).astype(np.float32) - return image diff --git a/taming/data/coco.py b/taming/data/coco.py deleted file mode 100644 index 2b2f783..0000000 --- a/taming/data/coco.py +++ /dev/null @@ -1,176 +0,0 @@ -import os -import json -import albumentations -import numpy as np -from PIL import Image -from tqdm import tqdm -from torch.utils.data import Dataset - -from taming.data.sflckr import SegmentationBase # for examples included in repo - - -class Examples(SegmentationBase): - def __init__(self, size=256, random_crop=False, interpolation="bicubic"): - super().__init__(data_csv="data/coco_examples.txt", - data_root="data/coco_images", - segmentation_root="data/coco_segmentations", - size=size, random_crop=random_crop, - interpolation=interpolation, - n_labels=183, shift_segmentation=True) - - -class CocoBase(Dataset): - """needed for (image, caption, segmentation) pairs""" - def __init__(self, size=None, dataroot="", datajson="", onehot_segmentation=False, use_stuffthing=False, - crop_size=None, force_no_crop=False, given_files=None): - self.split = self.get_split() - self.size = size - if crop_size is None: - self.crop_size = size - else: - self.crop_size = crop_size - - self.onehot = onehot_segmentation # return segmentation as rgb or one hot - self.stuffthing = use_stuffthing # include thing in segmentation - if self.onehot and not self.stuffthing: - raise NotImplemented("One hot mode is only supported for the " - "stuffthings version because labels are stored " - "a bit different.") - - data_json = datajson - with open(data_json) as json_file: - self.json_data = json.load(json_file) - self.img_id_to_captions = dict() - self.img_id_to_filepath = dict() - self.img_id_to_segmentation_filepath = dict() - - assert data_json.split("/")[-1] in ["captions_train2017.json", - "captions_val2017.json"] - if self.stuffthing: - self.segmentation_prefix = ( - "data/cocostuffthings/val2017" if - data_json.endswith("captions_val2017.json") else - "data/cocostuffthings/train2017") - else: - self.segmentation_prefix = ( - "data/coco/annotations/stuff_val2017_pixelmaps" if - data_json.endswith("captions_val2017.json") else - "data/coco/annotations/stuff_train2017_pixelmaps") - - imagedirs = self.json_data["images"] - self.labels = {"image_ids": list()} - for imgdir in tqdm(imagedirs, desc="ImgToPath"): - self.img_id_to_filepath[imgdir["id"]] = os.path.join(dataroot, imgdir["file_name"]) - self.img_id_to_captions[imgdir["id"]] = list() - pngfilename = imgdir["file_name"].replace("jpg", "png") - self.img_id_to_segmentation_filepath[imgdir["id"]] = os.path.join( - self.segmentation_prefix, pngfilename) - if given_files is not None: - if pngfilename in given_files: - self.labels["image_ids"].append(imgdir["id"]) - else: - self.labels["image_ids"].append(imgdir["id"]) - - capdirs = self.json_data["annotations"] - for capdir in tqdm(capdirs, desc="ImgToCaptions"): - # there are in average 5 captions per image - self.img_id_to_captions[capdir["image_id"]].append(np.array([capdir["caption"]])) - - self.rescaler = albumentations.SmallestMaxSize(max_size=self.size) - if self.split=="validation": - self.cropper = albumentations.CenterCrop(height=self.crop_size, width=self.crop_size) - else: - self.cropper = albumentations.RandomCrop(height=self.crop_size, width=self.crop_size) - self.preprocessor = albumentations.Compose( - [self.rescaler, self.cropper], - additional_targets={"segmentation": "image"}) - if force_no_crop: - self.rescaler = albumentations.Resize(height=self.size, width=self.size) - self.preprocessor = albumentations.Compose( - [self.rescaler], - additional_targets={"segmentation": "image"}) - - def __len__(self): - return len(self.labels["image_ids"]) - - def preprocess_image(self, image_path, segmentation_path): - image = Image.open(image_path) - if not image.mode == "RGB": - image = image.convert("RGB") - image = np.array(image).astype(np.uint8) - - segmentation = Image.open(segmentation_path) - if not self.onehot and not segmentation.mode == "RGB": - segmentation = segmentation.convert("RGB") - segmentation = np.array(segmentation).astype(np.uint8) - if self.onehot: - assert self.stuffthing - # stored in caffe format: unlabeled==255. stuff and thing from - # 0-181. to be compatible with the labels in - # https://github.com/nightrome/cocostuff/blob/master/labels.txt - # we shift stuffthing one to the right and put unlabeled in zero - # as long as segmentation is uint8 shifting to right handles the - # latter too - assert segmentation.dtype == np.uint8 - segmentation = segmentation + 1 - - processed = self.preprocessor(image=image, segmentation=segmentation) - image, segmentation = processed["image"], processed["segmentation"] - image = (image / 127.5 - 1.0).astype(np.float32) - - if self.onehot: - assert segmentation.dtype == np.uint8 - # make it one hot - n_labels = 183 - flatseg = np.ravel(segmentation) - onehot = np.zeros((flatseg.size, n_labels), dtype=np.bool) - onehot[np.arange(flatseg.size), flatseg] = True - onehot = onehot.reshape(segmentation.shape + (n_labels,)).astype(int) - segmentation = onehot - else: - segmentation = (segmentation / 127.5 - 1.0).astype(np.float32) - return image, segmentation - - def __getitem__(self, i): - img_path = self.img_id_to_filepath[self.labels["image_ids"][i]] - seg_path = self.img_id_to_segmentation_filepath[self.labels["image_ids"][i]] - image, segmentation = self.preprocess_image(img_path, seg_path) - captions = self.img_id_to_captions[self.labels["image_ids"][i]] - # randomly draw one of all available captions per image - caption = captions[np.random.randint(0, len(captions))] - example = {"image": image, - "caption": [str(caption[0])], - "segmentation": segmentation, - "img_path": img_path, - "seg_path": seg_path, - "filename_": img_path.split(os.sep)[-1] - } - return example - - -class CocoImagesAndCaptionsTrain(CocoBase): - """returns a pair of (image, caption)""" - def __init__(self, size, onehot_segmentation=False, use_stuffthing=False, crop_size=None, force_no_crop=False): - super().__init__(size=size, - dataroot="data/coco/train2017", - datajson="data/coco/annotations/captions_train2017.json", - onehot_segmentation=onehot_segmentation, - use_stuffthing=use_stuffthing, crop_size=crop_size, force_no_crop=force_no_crop) - - def get_split(self): - return "train" - - -class CocoImagesAndCaptionsValidation(CocoBase): - """returns a pair of (image, caption)""" - def __init__(self, size, onehot_segmentation=False, use_stuffthing=False, crop_size=None, force_no_crop=False, - given_files=None): - super().__init__(size=size, - dataroot="data/coco/val2017", - datajson="data/coco/annotations/captions_val2017.json", - onehot_segmentation=onehot_segmentation, - use_stuffthing=use_stuffthing, crop_size=crop_size, force_no_crop=force_no_crop, - given_files=given_files) - - def get_split(self): - return "validation" diff --git a/taming/data/conditional_builder/objects_bbox.py b/taming/data/conditional_builder/objects_bbox.py deleted file mode 100644 index 15881e7..0000000 --- a/taming/data/conditional_builder/objects_bbox.py +++ /dev/null @@ -1,60 +0,0 @@ -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 taming.data.image_transforms import convert_pil_to_tensor -from torch import LongTensor, Tensor - -from taming.data.helper_types import BoundingBox, Annotation -from taming.data.conditional_builder.objects_center_points import ObjectsCenterPointsConditionalBuilder -from taming.data.conditional_builder.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 - - def _make_object_descriptors(self, annotations: List[Annotation]) -> List[Tuple[int, ...]]: - object_triples = [ - (self.object_representation(ann), *self.token_pair_from_bbox(ann.bbox)) - for ann in annotations - ] - empty_triple = (self.none, self.none, self.none) - object_triples = pad_list(object_triples, empty_triple, self.no_max_objects) - return object_triples - - def inverse_build(self, conditional: LongTensor) -> Tuple[List[Tuple[int, BoundingBox]], Optional[BoundingBox]]: - conditional_list = conditional.tolist() - crop_coordinates = None - if self.encode_crop: - crop_coordinates = self.bbox_from_token_pair(conditional_list[-2], conditional_list[-1]) - conditional_list = conditional_list[:-2] - 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 - ], crop_coordinates - - 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) - ) - 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_no) + ' ' + additional_parameters_string(annotation) - 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/taming/data/conditional_builder/objects_center_points.py b/taming/data/conditional_builder/objects_center_points.py deleted file mode 100644 index 9a48032..0000000 --- a/taming/data/conditional_builder/objects_center_points.py +++ /dev/null @@ -1,168 +0,0 @@ -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 taming.data.conditional_builder.utils import COLOR_PALETTE, WHITE, GRAY_75, BLACK, FULL_CROP, filter_annotations, \ - additional_parameters_string, horizontally_flip_bbox, pad_list, get_circle_size, get_plot_font_size, \ - absolute_bbox, rescale_annotations -from taming.data.helper_types import BoundingBox, Annotation -from taming.data.image_transforms import convert_pil_to_tensor -from torch import LongTensor, Tensor - - -class ObjectsCenterPointsConditionalBuilder: - def __init__(self, no_object_classes: int, no_max_objects: int, no_tokens: int, encode_crop: bool, - use_group_parameter: bool, use_additional_parameters: bool): - self.no_object_classes = no_object_classes - self.no_max_objects = no_max_objects - self.no_tokens = no_tokens - self.encode_crop = encode_crop - self.no_sections = int(math.sqrt(self.no_tokens)) - self.use_group_parameter = use_group_parameter - self.use_additional_parameters = use_additional_parameters - - @property - def none(self) -> int: - return self.no_tokens - 1 - - @property - def object_descriptor_length(self) -> int: - return 2 - - @property - def embedding_dim(self) -> int: - extra_length = 2 if self.encode_crop else 0 - return self.no_max_objects * self.object_descriptor_length + extra_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 - 1))) - y_discrete = int(round(y * (self.no_sections - 1))) - return y_discrete * self.no_sections + x_discrete - - def coordinates_from_token(self, token: int) -> (float, float): - x = token % self.no_sections - y = token // self.no_sections - return x / (self.no_sections - 1), y / (self.no_sections - 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) - return x0, y0, x1 - x0, y1 - y0 - - def token_pair_from_bbox(self, bbox: BoundingBox) -> Tuple[int, int]: - return self.tokenize_coordinates(bbox[0], bbox[1]), \ - self.tokenize_coordinates(bbox[0] + bbox[2], bbox[1] + bbox[3]) - - def inverse_build(self, conditional: LongTensor) \ - -> Tuple[List[Tuple[int, Tuple[float, float]]], Optional[BoundingBox]]: - conditional_list = conditional.tolist() - crop_coordinates = None - if self.encode_crop: - crop_coordinates = self.bbox_from_token_pair(conditional_list[-2], conditional_list[-1]) - conditional_list = conditional_list[:-2] - 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 - ], crop_coordinates - - 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)) - 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_no) + ' ' + 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: - modifier = 0 - if self.use_group_parameter: - modifier |= 1 * (annotation.is_group_of is True) - if self.use_additional_parameters: - modifier |= 2 * (annotation.is_occluded is True) - modifier |= 4 * (annotation.is_depiction is True) - modifier |= 8 * (annotation.is_inside is True) - return annotation.category_no + self.no_object_classes * modifier - - def representation_to_annotation(self, representation: int) -> Annotation: - category_no = representation % self.no_object_classes - modifier = representation // self.no_object_classes - # noinspection PyTypeChecker - return Annotation( - area=None, image_id=None, bbox=None, category_id=None, id=None, source=None, confidence=None, - category_no=category_no, - is_group_of=bool((modifier & 1) * self.use_group_parameter), - is_occluded=bool((modifier & 2) * self.use_additional_parameters), - is_depiction=bool((modifier & 4) * self.use_additional_parameters), - is_inside=bool((modifier & 8) * self.use_additional_parameters) - ) - - def _crop_encoder(self, crop_coordinates: BoundingBox) -> List[int]: - return list(self.token_pair_from_bbox(crop_coordinates)) - - def _make_object_descriptors(self, annotations: List[Annotation]) -> List[Tuple[int, ...]]: - object_tuples = [ - (self.object_representation(a), - self.tokenize_coordinates(a.bbox[0] + a.bbox[2] / 2, a.bbox[1] + a.bbox[3] / 2)) - 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, crop_coordinates: Optional[BoundingBox] = None, horizontal_flip: bool = False) \ - -> LongTensor: - if len(annotations) == 0: - warnings.warn('Did not receive any annotations.') - if len(annotations) > self.no_max_objects: - warnings.warn('Received more annotations than allowed.') - annotations = annotations[:self.no_max_objects] - - if not crop_coordinates: - crop_coordinates = FULL_CROP - - random.shuffle(annotations) - annotations = filter_annotations(annotations, crop_coordinates) - if self.encode_crop: - annotations = rescale_annotations(annotations, FULL_CROP, horizontal_flip) - if horizontal_flip: - crop_coordinates = horizontally_flip_bbox(crop_coordinates) - extra = self._crop_encoder(crop_coordinates) - else: - annotations = rescale_annotations(annotations, crop_coordinates, horizontal_flip) - extra = [] - - object_tuples = self._make_object_descriptors(annotations) - flattened = [token for tuple_ in object_tuples for token in tuple_] + extra - assert len(flattened) == self.embedding_dim - assert all(0 <= value < self.no_tokens for value in flattened) - return LongTensor(flattened) diff --git a/taming/data/conditional_builder/utils.py b/taming/data/conditional_builder/utils.py deleted file mode 100644 index d0ee175..0000000 --- a/taming/data/conditional_builder/utils.py +++ /dev/null @@ -1,105 +0,0 @@ -import importlib -from typing import List, Any, Tuple, Optional - -from taming.data.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 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 - return int(bbox[0]), int(bbox[1]), int(bbox[2]), int(bbox[3]) - - -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/taming/data/custom.py b/taming/data/custom.py deleted file mode 100644 index 33f302a..0000000 --- a/taming/data/custom.py +++ /dev/null @@ -1,38 +0,0 @@ -import os -import numpy as np -import albumentations -from torch.utils.data import Dataset - -from taming.data.base import ImagePaths, NumpyPaths, ConcatDatasetWithIndex - - -class CustomBase(Dataset): - def __init__(self, *args, **kwargs): - super().__init__() - self.data = None - - def __len__(self): - return len(self.data) - - def __getitem__(self, i): - example = self.data[i] - return example - - - -class CustomTrain(CustomBase): - def __init__(self, size, training_images_list_file): - super().__init__() - with open(training_images_list_file, "r") as f: - paths = f.read().splitlines() - self.data = ImagePaths(paths=paths, size=size, random_crop=False) - - -class CustomTest(CustomBase): - def __init__(self, size, test_images_list_file): - super().__init__() - with open(test_images_list_file, "r") as f: - paths = f.read().splitlines() - self.data = ImagePaths(paths=paths, size=size, random_crop=False) - - diff --git a/taming/data/faceshq.py b/taming/data/faceshq.py deleted file mode 100644 index 6912d04..0000000 --- a/taming/data/faceshq.py +++ /dev/null @@ -1,134 +0,0 @@ -import os -import numpy as np -import albumentations -from torch.utils.data import Dataset - -from taming.data.base import ImagePaths, NumpyPaths, ConcatDatasetWithIndex - - -class FacesBase(Dataset): - def __init__(self, *args, **kwargs): - super().__init__() - self.data = None - self.keys = None - - def __len__(self): - return len(self.data) - - def __getitem__(self, i): - example = self.data[i] - ex = {} - if self.keys is not None: - for k in self.keys: - ex[k] = example[k] - else: - ex = example - return ex - - -class CelebAHQTrain(FacesBase): - def __init__(self, size, keys=None): - super().__init__() - root = "data/celebahq" - with open("data/celebahqtrain.txt", "r") as f: - relpaths = f.read().splitlines() - paths = [os.path.join(root, relpath) for relpath in relpaths] - self.data = NumpyPaths(paths=paths, size=size, random_crop=False) - self.keys = keys - - -class CelebAHQValidation(FacesBase): - def __init__(self, size, keys=None): - super().__init__() - root = "data/celebahq" - with open("data/celebahqvalidation.txt", "r") as f: - relpaths = f.read().splitlines() - paths = [os.path.join(root, relpath) for relpath in relpaths] - self.data = NumpyPaths(paths=paths, size=size, random_crop=False) - self.keys = keys - - -class FFHQTrain(FacesBase): - def __init__(self, size, keys=None): - super().__init__() - root = "data/ffhq" - with open("data/ffhqtrain.txt", "r") as f: - relpaths = f.read().splitlines() - paths = [os.path.join(root, relpath) for relpath in relpaths] - self.data = ImagePaths(paths=paths, size=size, random_crop=False) - self.keys = keys - - -class FFHQValidation(FacesBase): - def __init__(self, size, keys=None): - super().__init__() - root = "data/ffhq" - with open("data/ffhqvalidation.txt", "r") as f: - relpaths = f.read().splitlines() - paths = [os.path.join(root, relpath) for relpath in relpaths] - self.data = ImagePaths(paths=paths, size=size, random_crop=False) - self.keys = keys - - -class FacesHQTrain(Dataset): - # CelebAHQ [0] + FFHQ [1] - def __init__(self, size, keys=None, crop_size=None, coord=False): - d1 = CelebAHQTrain(size=size, keys=keys) - d2 = FFHQTrain(size=size, keys=keys) - self.data = ConcatDatasetWithIndex([d1, d2]) - self.coord = coord - if crop_size is not None: - self.cropper = albumentations.RandomCrop(height=crop_size,width=crop_size) - if self.coord: - self.cropper = albumentations.Compose([self.cropper], - additional_targets={"coord": "image"}) - - def __len__(self): - return len(self.data) - - def __getitem__(self, i): - ex, y = self.data[i] - if hasattr(self, "cropper"): - if not self.coord: - out = self.cropper(image=ex["image"]) - ex["image"] = out["image"] - else: - h,w,_ = ex["image"].shape - coord = np.arange(h*w).reshape(h,w,1)/(h*w) - out = self.cropper(image=ex["image"], coord=coord) - ex["image"] = out["image"] - ex["coord"] = out["coord"] - ex["class"] = y - return ex - - -class FacesHQValidation(Dataset): - # CelebAHQ [0] + FFHQ [1] - def __init__(self, size, keys=None, crop_size=None, coord=False): - d1 = CelebAHQValidation(size=size, keys=keys) - d2 = FFHQValidation(size=size, keys=keys) - self.data = ConcatDatasetWithIndex([d1, d2]) - self.coord = coord - if crop_size is not None: - self.cropper = albumentations.CenterCrop(height=crop_size,width=crop_size) - if self.coord: - self.cropper = albumentations.Compose([self.cropper], - additional_targets={"coord": "image"}) - - def __len__(self): - return len(self.data) - - def __getitem__(self, i): - ex, y = self.data[i] - if hasattr(self, "cropper"): - if not self.coord: - out = self.cropper(image=ex["image"]) - ex["image"] = out["image"] - else: - h,w,_ = ex["image"].shape - coord = np.arange(h*w).reshape(h,w,1)/(h*w) - out = self.cropper(image=ex["image"], coord=coord) - ex["image"] = out["image"] - ex["coord"] = out["coord"] - ex["class"] = y - return ex diff --git a/taming/data/helper_types.py b/taming/data/helper_types.py deleted file mode 100644 index fb51e30..0000000 --- a/taming/data/helper_types.py +++ /dev/null @@ -1,49 +0,0 @@ -from typing import Dict, Tuple, Optional, NamedTuple, Union -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 -CropMethodType = Literal['none', 'random', 'center', 'random-2d'] -SplitType = Literal['train', 'validation', 'test'] - - -class ImageDescription(NamedTuple): - id: int - file_name: str - original_size: Tuple[int, int] # w, h - url: Optional[str] = None - license: Optional[int] = None - coco_url: Optional[str] = None - date_captured: Optional[str] = None - flickr_url: Optional[str] = None - flickr_id: Optional[str] = None - coco_id: Optional[str] = None - - -class Category(NamedTuple): - id: str - super_category: Optional[str] - name: str - - -class Annotation(NamedTuple): - area: float - image_id: str - bbox: BoundingBox - category_no: int - category_id: str - id: Optional[int] = None - source: Optional[str] = None - confidence: Optional[float] = None - is_group_of: Optional[bool] = None - is_truncated: Optional[bool] = None - is_occluded: Optional[bool] = None - is_depiction: Optional[bool] = None - is_inside: Optional[bool] = None - segmentation: Optional[Dict] = None diff --git a/taming/data/image_transforms.py b/taming/data/image_transforms.py deleted file mode 100644 index 657ac33..0000000 --- a/taming/data/image_transforms.py +++ /dev/null @@ -1,132 +0,0 @@ -import random -import warnings -from typing import Union - -import torch -from torch import Tensor -from torchvision.transforms import RandomCrop, functional as F, CenterCrop, RandomHorizontalFlip, PILToTensor -from torchvision.transforms.functional import _get_image_size as get_image_size - -from taming.data.helper_types import BoundingBox, Image - -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 RandomCrop1dReturnCoordinates(RandomCrop): - def forward(self, img: Image) -> (BoundingBox, Image): - """ - Additionally to cropping, returns the relative coordinates of the crop bounding box. - Args: - img (PIL Image or Tensor): Image to be cropped. - - Returns: - Bounding box: x0, y0, w, h - PIL Image or Tensor: Cropped image. - - Based on: - torchvision.transforms.RandomCrop, torchvision 1.7.0 - """ - if self.padding is not None: - img = F.pad(img, self.padding, self.fill, self.padding_mode) - - width, height = get_image_size(img) - # pad the width if needed - if self.pad_if_needed and width < self.size[1]: - padding = [self.size[1] - width, 0] - img = F.pad(img, padding, self.fill, self.padding_mode) - # pad the height if needed - if self.pad_if_needed and height < self.size[0]: - padding = [0, self.size[0] - height] - img = F.pad(img, padding, self.fill, self.padding_mode) - - i, j, h, w = self.get_params(img, self.size) - bbox = (j / width, i / height, w / width, h / height) # x0, y0, w, h - return bbox, F.crop(img, i, j, h, w) - - -class Random2dCropReturnCoordinates(torch.nn.Module): - """ - Additionally to cropping, returns the relative coordinates of the crop bounding box. - Args: - img (PIL Image or Tensor): Image to be cropped. - - Returns: - Bounding box: x0, y0, w, h - PIL Image or Tensor: Cropped image. - - Based on: - torchvision.transforms.RandomCrop, torchvision 1.7.0 - """ - - def __init__(self, min_size: int): - super().__init__() - self.min_size = min_size - - def forward(self, img: Image) -> (BoundingBox, Image): - width, height = get_image_size(img) - max_size = min(width, height) - if max_size <= self.min_size: - size = max_size - else: - size = random.randint(self.min_size, max_size) - top = random.randint(0, height - size) - left = random.randint(0, width - size) - bbox = left / width, top / height, size / width, size / height - return bbox, F.crop(img, top, left, size, size) - - -class CenterCropReturnCoordinates(CenterCrop): - @staticmethod - def get_bbox_of_center_crop(width: int, height: int) -> BoundingBox: - if width > height: - w = height / width - h = 1.0 - x0 = 0.5 - w / 2 - y0 = 0. - else: - w = 1.0 - h = width / height - x0 = 0. - y0 = 0.5 - h / 2 - return x0, y0, w, h - - def forward(self, img: Union[Image, Tensor]) -> (BoundingBox, Union[Image, Tensor]): - """ - Additionally to cropping, returns the relative coordinates of the crop bounding box. - Args: - img (PIL Image or Tensor): Image to be cropped. - - Returns: - Bounding box: x0, y0, w, h - PIL Image or Tensor: Cropped image. - Based on: - torchvision.transforms.RandomHorizontalFlip (version 1.7.0) - """ - width, height = get_image_size(img) - return self.get_bbox_of_center_crop(width, height), F.center_crop(img, self.size) - - -class RandomHorizontalFlipReturn(RandomHorizontalFlip): - def forward(self, img: Image) -> (bool, Image): - """ - Additionally to flipping, returns a boolean whether it was flipped or not. - Args: - img (PIL Image or Tensor): Image to be flipped. - - Returns: - flipped: whether the image was flipped or not - PIL Image or Tensor: Randomly flipped image. - - Based on: - torchvision.transforms.RandomHorizontalFlip (version 1.7.0) - """ - if torch.rand(1) < self.p: - return True, F.hflip(img) - return False, img diff --git a/taming/data/imagenet.py b/taming/data/imagenet.py deleted file mode 100644 index 9a02ec4..0000000 --- a/taming/data/imagenet.py +++ /dev/null @@ -1,558 +0,0 @@ -import os, tarfile, glob, shutil -import yaml -import numpy as np -from tqdm import tqdm -from PIL import Image -import albumentations -from omegaconf import OmegaConf -from torch.utils.data import Dataset - -from taming.data.base import ImagePaths -from taming.util import download, retrieve -import taming.data.utils as bdu - - -def give_synsets_from_indices(indices, path_to_yaml="data/imagenet_idx_to_synset.yaml"): - synsets = [] - with open(path_to_yaml) as f: - di2s = yaml.load(f) - for idx in indices: - synsets.append(str(di2s[idx])) - print("Using {} different synsets for construction of Restriced Imagenet.".format(len(synsets))) - return synsets - - -def str_to_indices(string): - """Expects a string in the format '32-123, 256, 280-321'""" - assert not string.endswith(","), "provided string '{}' ends with a comma, pls remove it".format(string) - subs = string.split(",") - indices = [] - for sub in subs: - subsubs = sub.split("-") - assert len(subsubs) > 0 - if len(subsubs) == 1: - indices.append(int(subsubs[0])) - else: - rang = [j for j in range(int(subsubs[0]), int(subsubs[1]))] - indices.extend(rang) - return sorted(indices) - - -class ImageNetBase(Dataset): - def __init__(self, config=None): - self.config = config or OmegaConf.create() - if not type(self.config)==dict: - self.config = OmegaConf.to_container(self.config) - self._prepare() - self._prepare_synset_to_human() - self._prepare_idx_to_synset() - self._load() - - def __len__(self): - return len(self.data) - - def __getitem__(self, i): - return self.data[i] - - def _prepare(self): - raise NotImplementedError() - - def _filter_relpaths(self, relpaths): - ignore = set([ - "n06596364_9591.JPEG", - ]) - relpaths = [rpath for rpath in relpaths if not rpath.split("/")[-1] in ignore] - if "sub_indices" in self.config: - indices = str_to_indices(self.config["sub_indices"]) - synsets = give_synsets_from_indices(indices, path_to_yaml=self.idx2syn) # returns a list of strings - files = [] - for rpath in relpaths: - syn = rpath.split("/")[0] - if syn in synsets: - files.append(rpath) - return files - else: - return relpaths - - def _prepare_synset_to_human(self): - SIZE = 2655750 - URL = "https://heibox.uni-heidelberg.de/f/9f28e956cd304264bb82/?dl=1" - self.human_dict = os.path.join(self.root, "synset_human.txt") - if (not os.path.exists(self.human_dict) or - not os.path.getsize(self.human_dict)==SIZE): - download(URL, self.human_dict) - - def _prepare_idx_to_synset(self): - URL = "https://heibox.uni-heidelberg.de/f/d835d5b6ceda4d3aa910/?dl=1" - self.idx2syn = os.path.join(self.root, "index_synset.yaml") - if (not os.path.exists(self.idx2syn)): - download(URL, self.idx2syn) - - def _load(self): - with open(self.txt_filelist, "r") as f: - self.relpaths = f.read().splitlines() - l1 = len(self.relpaths) - self.relpaths = self._filter_relpaths(self.relpaths) - print("Removed {} files from filelist during filtering.".format(l1 - len(self.relpaths))) - - self.synsets = [p.split("/")[0] for p in self.relpaths] - self.abspaths = [os.path.join(self.datadir, p) for p in self.relpaths] - - unique_synsets = np.unique(self.synsets) - class_dict = dict((synset, i) for i, synset in enumerate(unique_synsets)) - self.class_labels = [class_dict[s] for s in self.synsets] - - with open(self.human_dict, "r") as f: - human_dict = f.read().splitlines() - human_dict = dict(line.split(maxsplit=1) for line in human_dict) - - self.human_labels = [human_dict[s] for s in self.synsets] - - labels = { - "relpath": np.array(self.relpaths), - "synsets": np.array(self.synsets), - "class_label": np.array(self.class_labels), - "human_label": np.array(self.human_labels), - } - self.data = ImagePaths(self.abspaths, - labels=labels, - size=retrieve(self.config, "size", default=0), - random_crop=self.random_crop) - - -class ImageNetTrain(ImageNetBase): - NAME = "ILSVRC2012_train" - URL = "http://www.image-net.org/challenges/LSVRC/2012/" - AT_HASH = "a306397ccf9c2ead27155983c254227c0fd938e2" - FILES = [ - "ILSVRC2012_img_train.tar", - ] - SIZES = [ - 147897477120, - ] - - def _prepare(self): - self.random_crop = retrieve(self.config, "ImageNetTrain/random_crop", - default=True) - cachedir = os.environ.get("XDG_CACHE_HOME", os.path.expanduser("~/.cache")) - self.root = os.path.join(cachedir, "autoencoders/data", self.NAME) - self.datadir = os.path.join(self.root, "data") - self.txt_filelist = os.path.join(self.root, "filelist.txt") - self.expected_length = 1281167 - if not bdu.is_prepared(self.root): - # prep - print("Preparing dataset {} in {}".format(self.NAME, self.root)) - - datadir = self.datadir - if not os.path.exists(datadir): - path = os.path.join(self.root, self.FILES[0]) - if not os.path.exists(path) or not os.path.getsize(path)==self.SIZES[0]: - import academictorrents as at - atpath = at.get(self.AT_HASH, datastore=self.root) - assert atpath == path - - print("Extracting {} to {}".format(path, datadir)) - os.makedirs(datadir, exist_ok=True) - with tarfile.open(path, "r:") as tar: - tar.extractall(path=datadir) - - print("Extracting sub-tars.") - subpaths = sorted(glob.glob(os.path.join(datadir, "*.tar"))) - for subpath in tqdm(subpaths): - subdir = subpath[:-len(".tar")] - os.makedirs(subdir, exist_ok=True) - with tarfile.open(subpath, "r:") as tar: - tar.extractall(path=subdir) - - - filelist = glob.glob(os.path.join(datadir, "**", "*.JPEG")) - filelist = [os.path.relpath(p, start=datadir) for p in filelist] - filelist = sorted(filelist) - filelist = "\n".join(filelist)+"\n" - with open(self.txt_filelist, "w") as f: - f.write(filelist) - - bdu.mark_prepared(self.root) - - -class ImageNetValidation(ImageNetBase): - NAME = "ILSVRC2012_validation" - URL = "http://www.image-net.org/challenges/LSVRC/2012/" - AT_HASH = "5d6d0df7ed81efd49ca99ea4737e0ae5e3a5f2e5" - VS_URL = "https://heibox.uni-heidelberg.de/f/3e0f6e9c624e45f2bd73/?dl=1" - FILES = [ - "ILSVRC2012_img_val.tar", - "validation_synset.txt", - ] - SIZES = [ - 6744924160, - 1950000, - ] - - def _prepare(self): - self.random_crop = retrieve(self.config, "ImageNetValidation/random_crop", - default=False) - cachedir = os.environ.get("XDG_CACHE_HOME", os.path.expanduser("~/.cache")) - self.root = os.path.join(cachedir, "autoencoders/data", self.NAME) - self.datadir = os.path.join(self.root, "data") - self.txt_filelist = os.path.join(self.root, "filelist.txt") - self.expected_length = 50000 - if not bdu.is_prepared(self.root): - # prep - print("Preparing dataset {} in {}".format(self.NAME, self.root)) - - datadir = self.datadir - if not os.path.exists(datadir): - path = os.path.join(self.root, self.FILES[0]) - if not os.path.exists(path) or not os.path.getsize(path)==self.SIZES[0]: - import academictorrents as at - atpath = at.get(self.AT_HASH, datastore=self.root) - assert atpath == path - - print("Extracting {} to {}".format(path, datadir)) - os.makedirs(datadir, exist_ok=True) - with tarfile.open(path, "r:") as tar: - tar.extractall(path=datadir) - - vspath = os.path.join(self.root, self.FILES[1]) - if not os.path.exists(vspath) or not os.path.getsize(vspath)==self.SIZES[1]: - download(self.VS_URL, vspath) - - with open(vspath, "r") as f: - synset_dict = f.read().splitlines() - synset_dict = dict(line.split() for line in synset_dict) - - print("Reorganizing into synset folders") - synsets = np.unique(list(synset_dict.values())) - for s in synsets: - os.makedirs(os.path.join(datadir, s), exist_ok=True) - for k, v in synset_dict.items(): - src = os.path.join(datadir, k) - dst = os.path.join(datadir, v) - shutil.move(src, dst) - - filelist = glob.glob(os.path.join(datadir, "**", "*.JPEG")) - filelist = [os.path.relpath(p, start=datadir) for p in filelist] - filelist = sorted(filelist) - filelist = "\n".join(filelist)+"\n" - with open(self.txt_filelist, "w") as f: - f.write(filelist) - - bdu.mark_prepared(self.root) - - -def get_preprocessor(size=None, random_crop=False, additional_targets=None, - crop_size=None): - if size is not None and size > 0: - transforms = list() - rescaler = albumentations.SmallestMaxSize(max_size = size) - transforms.append(rescaler) - if not random_crop: - cropper = albumentations.CenterCrop(height=size,width=size) - transforms.append(cropper) - else: - cropper = albumentations.RandomCrop(height=size,width=size) - transforms.append(cropper) - flipper = albumentations.HorizontalFlip() - transforms.append(flipper) - preprocessor = albumentations.Compose(transforms, - additional_targets=additional_targets) - elif crop_size is not None and crop_size > 0: - if not random_crop: - cropper = albumentations.CenterCrop(height=crop_size,width=crop_size) - else: - cropper = albumentations.RandomCrop(height=crop_size,width=crop_size) - transforms = [cropper] - preprocessor = albumentations.Compose(transforms, - additional_targets=additional_targets) - else: - preprocessor = lambda **kwargs: kwargs - return preprocessor - - -def rgba_to_depth(x): - assert x.dtype == np.uint8 - assert len(x.shape) == 3 and x.shape[2] == 4 - y = x.copy() - y.dtype = np.float32 - y = y.reshape(x.shape[:2]) - return np.ascontiguousarray(y) - - -class BaseWithDepth(Dataset): - DEFAULT_DEPTH_ROOT="data/imagenet_depth" - - def __init__(self, config=None, size=None, random_crop=False, - crop_size=None, root=None): - self.config = config - self.base_dset = self.get_base_dset() - self.preprocessor = get_preprocessor( - size=size, - crop_size=crop_size, - random_crop=random_crop, - additional_targets={"depth": "image"}) - self.crop_size = crop_size - if self.crop_size is not None: - self.rescaler = albumentations.Compose( - [albumentations.SmallestMaxSize(max_size = self.crop_size)], - additional_targets={"depth": "image"}) - if root is not None: - self.DEFAULT_DEPTH_ROOT = root - - def __len__(self): - return len(self.base_dset) - - def preprocess_depth(self, path): - rgba = np.array(Image.open(path)) - depth = rgba_to_depth(rgba) - depth = (depth - depth.min())/max(1e-8, depth.max()-depth.min()) - depth = 2.0*depth-1.0 - return depth - - def __getitem__(self, i): - e = self.base_dset[i] - e["depth"] = self.preprocess_depth(self.get_depth_path(e)) - # up if necessary - h,w,c = e["image"].shape - if self.crop_size and min(h,w) < self.crop_size: - # have to upscale to be able to crop - this just uses bilinear - out = self.rescaler(image=e["image"], depth=e["depth"]) - e["image"] = out["image"] - e["depth"] = out["depth"] - transformed = self.preprocessor(image=e["image"], depth=e["depth"]) - e["image"] = transformed["image"] - e["depth"] = transformed["depth"] - return e - - -class ImageNetTrainWithDepth(BaseWithDepth): - # default to random_crop=True - def __init__(self, random_crop=True, sub_indices=None, **kwargs): - self.sub_indices = sub_indices - super().__init__(random_crop=random_crop, **kwargs) - - def get_base_dset(self): - if self.sub_indices is None: - return ImageNetTrain() - else: - return ImageNetTrain({"sub_indices": self.sub_indices}) - - def get_depth_path(self, e): - fid = os.path.splitext(e["relpath"])[0]+".png" - fid = os.path.join(self.DEFAULT_DEPTH_ROOT, "train", fid) - return fid - - -class ImageNetValidationWithDepth(BaseWithDepth): - def __init__(self, sub_indices=None, **kwargs): - self.sub_indices = sub_indices - super().__init__(**kwargs) - - def get_base_dset(self): - if self.sub_indices is None: - return ImageNetValidation() - else: - return ImageNetValidation({"sub_indices": self.sub_indices}) - - def get_depth_path(self, e): - fid = os.path.splitext(e["relpath"])[0]+".png" - fid = os.path.join(self.DEFAULT_DEPTH_ROOT, "val", fid) - return fid - - -class RINTrainWithDepth(ImageNetTrainWithDepth): - def __init__(self, config=None, size=None, random_crop=True, crop_size=None): - sub_indices = "30-32, 33-37, 151-268, 281-285, 80-100, 365-382, 389-397, 118-121, 300-319" - super().__init__(config=config, size=size, random_crop=random_crop, - sub_indices=sub_indices, crop_size=crop_size) - - -class RINValidationWithDepth(ImageNetValidationWithDepth): - def __init__(self, config=None, size=None, random_crop=False, crop_size=None): - sub_indices = "30-32, 33-37, 151-268, 281-285, 80-100, 365-382, 389-397, 118-121, 300-319" - super().__init__(config=config, size=size, random_crop=random_crop, - sub_indices=sub_indices, crop_size=crop_size) - - -class DRINExamples(Dataset): - def __init__(self): - self.preprocessor = get_preprocessor(size=256, additional_targets={"depth": "image"}) - with open("data/drin_examples.txt", "r") as f: - relpaths = f.read().splitlines() - self.image_paths = [os.path.join("data/drin_images", - relpath) for relpath in relpaths] - self.depth_paths = [os.path.join("data/drin_depth", - relpath.replace(".JPEG", ".png")) for relpath in relpaths] - - def __len__(self): - return len(self.image_paths) - - def preprocess_image(self, image_path): - image = Image.open(image_path) - if not image.mode == "RGB": - image = image.convert("RGB") - image = np.array(image).astype(np.uint8) - image = self.preprocessor(image=image)["image"] - image = (image/127.5 - 1.0).astype(np.float32) - return image - - def preprocess_depth(self, path): - rgba = np.array(Image.open(path)) - depth = rgba_to_depth(rgba) - depth = (depth - depth.min())/max(1e-8, depth.max()-depth.min()) - depth = 2.0*depth-1.0 - return depth - - def __getitem__(self, i): - e = dict() - e["image"] = self.preprocess_image(self.image_paths[i]) - e["depth"] = self.preprocess_depth(self.depth_paths[i]) - transformed = self.preprocessor(image=e["image"], depth=e["depth"]) - e["image"] = transformed["image"] - e["depth"] = transformed["depth"] - return e - - -def imscale(x, factor, keepshapes=False, keepmode="bicubic"): - if factor is None or factor==1: - return x - - dtype = x.dtype - assert dtype in [np.float32, np.float64] - assert x.min() >= -1 - assert x.max() <= 1 - - keepmode = {"nearest": Image.NEAREST, "bilinear": Image.BILINEAR, - "bicubic": Image.BICUBIC}[keepmode] - - lr = (x+1.0)*127.5 - lr = lr.clip(0,255).astype(np.uint8) - lr = Image.fromarray(lr) - - h, w, _ = x.shape - nh = h//factor - nw = w//factor - assert nh > 0 and nw > 0, (nh, nw) - - lr = lr.resize((nw,nh), Image.BICUBIC) - if keepshapes: - lr = lr.resize((w,h), keepmode) - lr = np.array(lr)/127.5-1.0 - lr = lr.astype(dtype) - - return lr - - -class ImageNetScale(Dataset): - def __init__(self, size=None, crop_size=None, random_crop=False, - up_factor=None, hr_factor=None, keep_mode="bicubic"): - self.base = self.get_base() - - self.size = size - self.crop_size = crop_size if crop_size is not None else self.size - self.random_crop = random_crop - self.up_factor = up_factor - self.hr_factor = hr_factor - self.keep_mode = keep_mode - - transforms = list() - - if self.size is not None and self.size > 0: - rescaler = albumentations.SmallestMaxSize(max_size = self.size) - self.rescaler = rescaler - transforms.append(rescaler) - - if self.crop_size is not None and self.crop_size > 0: - if len(transforms) == 0: - self.rescaler = albumentations.SmallestMaxSize(max_size = self.crop_size) - - if not self.random_crop: - cropper = albumentations.CenterCrop(height=self.crop_size,width=self.crop_size) - else: - cropper = albumentations.RandomCrop(height=self.crop_size,width=self.crop_size) - transforms.append(cropper) - - if len(transforms) > 0: - if self.up_factor is not None: - additional_targets = {"lr": "image"} - else: - additional_targets = None - self.preprocessor = albumentations.Compose(transforms, - additional_targets=additional_targets) - else: - self.preprocessor = lambda **kwargs: kwargs - - def __len__(self): - return len(self.base) - - def __getitem__(self, i): - example = self.base[i] - image = example["image"] - # adjust resolution - image = imscale(image, self.hr_factor, keepshapes=False) - h,w,c = image.shape - if self.crop_size and min(h,w) < self.crop_size: - # have to upscale to be able to crop - this just uses bilinear - image = self.rescaler(image=image)["image"] - if self.up_factor is None: - image = self.preprocessor(image=image)["image"] - example["image"] = image - else: - lr = imscale(image, self.up_factor, keepshapes=True, - keepmode=self.keep_mode) - - out = self.preprocessor(image=image, lr=lr) - example["image"] = out["image"] - example["lr"] = out["lr"] - - return example - -class ImageNetScaleTrain(ImageNetScale): - def __init__(self, random_crop=True, **kwargs): - super().__init__(random_crop=random_crop, **kwargs) - - def get_base(self): - return ImageNetTrain() - -class ImageNetScaleValidation(ImageNetScale): - def get_base(self): - return ImageNetValidation() - - -from skimage.feature import canny -from skimage.color import rgb2gray - - -class ImageNetEdges(ImageNetScale): - def __init__(self, up_factor=1, **kwargs): - super().__init__(up_factor=1, **kwargs) - - def __getitem__(self, i): - example = self.base[i] - image = example["image"] - h,w,c = image.shape - if self.crop_size and min(h,w) < self.crop_size: - # have to upscale to be able to crop - this just uses bilinear - image = self.rescaler(image=image)["image"] - - lr = canny(rgb2gray(image), sigma=2) - lr = lr.astype(np.float32) - lr = lr[:,:,None][:,:,[0,0,0]] - - out = self.preprocessor(image=image, lr=lr) - example["image"] = out["image"] - example["lr"] = out["lr"] - - return example - - -class ImageNetEdgesTrain(ImageNetEdges): - def __init__(self, random_crop=True, **kwargs): - super().__init__(random_crop=random_crop, **kwargs) - - def get_base(self): - return ImageNetTrain() - -class ImageNetEdgesValidation(ImageNetEdges): - def get_base(self): - return ImageNetValidation() diff --git a/taming/data/open_images_helper.py b/taming/data/open_images_helper.py deleted file mode 100644 index 8feb7c6..0000000 --- a/taming/data/open_images_helper.py +++ /dev/null @@ -1,379 +0,0 @@ -open_images_unify_categories_for_coco = { - '/m/03bt1vf': '/m/01g317', - '/m/04yx4': '/m/01g317', - '/m/05r655': '/m/01g317', - '/m/01bl7v': '/m/01g317', - '/m/0cnyhnx': '/m/01xq0k1', - '/m/01226z': '/m/018xm', - '/m/05ctyq': '/m/018xm', - '/m/058qzx': '/m/04ctx', - '/m/06pcq': '/m/0l515', - '/m/03m3pdh': '/m/02crq1', - '/m/046dlr': '/m/01x3z', - '/m/0h8mzrc': '/m/01x3z', -} - - -top_300_classes_plus_coco_compatibility = [ - ('Man', 1060962), - ('Clothing', 986610), - ('Tree', 748162), - ('Woman', 611896), - ('Person', 610294), - ('Human face', 442948), - ('Girl', 175399), - ('Building', 162147), - ('Car', 159135), - ('Plant', 155704), - ('Human body', 137073), - ('Flower', 133128), - ('Window', 127485), - ('Human arm', 118380), - ('House', 114365), - ('Wheel', 111684), - ('Suit', 99054), - ('Human hair', 98089), - ('Human head', 92763), - ('Chair', 88624), - ('Boy', 79849), - ('Table', 73699), - ('Jeans', 57200), - ('Tire', 55725), - ('Skyscraper', 53321), - ('Food', 52400), - ('Footwear', 50335), - ('Dress', 50236), - ('Human leg', 47124), - ('Toy', 46636), - ('Tower', 45605), - ('Boat', 43486), - ('Land vehicle', 40541), - ('Bicycle wheel', 34646), - ('Palm tree', 33729), - ('Fashion accessory', 32914), - ('Glasses', 31940), - ('Bicycle', 31409), - ('Furniture', 30656), - ('Sculpture', 29643), - ('Bottle', 27558), - ('Dog', 26980), - ('Snack', 26796), - ('Human hand', 26664), - ('Bird', 25791), - ('Book', 25415), - ('Guitar', 24386), - ('Jacket', 23998), - ('Poster', 22192), - ('Dessert', 21284), - ('Baked goods', 20657), - ('Drink', 19754), - ('Flag', 18588), - ('Houseplant', 18205), - ('Tableware', 17613), - ('Airplane', 17218), - ('Door', 17195), - ('Sports uniform', 17068), - ('Shelf', 16865), - ('Drum', 16612), - ('Vehicle', 16542), - ('Microphone', 15269), - ('Street light', 14957), - ('Cat', 14879), - ('Fruit', 13684), - ('Fast food', 13536), - ('Animal', 12932), - ('Vegetable', 12534), - ('Train', 12358), - ('Horse', 11948), - ('Flowerpot', 11728), - ('Motorcycle', 11621), - ('Fish', 11517), - ('Desk', 11405), - ('Helmet', 10996), - ('Truck', 10915), - ('Bus', 10695), - ('Hat', 10532), - ('Auto part', 10488), - ('Musical instrument', 10303), - ('Sunglasses', 10207), - ('Picture frame', 10096), - ('Sports equipment', 10015), - ('Shorts', 9999), - ('Wine glass', 9632), - ('Duck', 9242), - ('Wine', 9032), - ('Rose', 8781), - ('Tie', 8693), - ('Butterfly', 8436), - ('Beer', 7978), - ('Cabinetry', 7956), - ('Laptop', 7907), - ('Insect', 7497), - ('Goggles', 7363), - ('Shirt', 7098), - ('Dairy Product', 7021), - ('Marine invertebrates', 7014), - ('Cattle', 7006), - ('Trousers', 6903), - ('Van', 6843), - ('Billboard', 6777), - ('Balloon', 6367), - ('Human nose', 6103), - ('Tent', 6073), - ('Camera', 6014), - ('Doll', 6002), - ('Coat', 5951), - ('Mobile phone', 5758), - ('Swimwear', 5729), - ('Strawberry', 5691), - ('Stairs', 5643), - ('Goose', 5599), - ('Umbrella', 5536), - ('Cake', 5508), - ('Sun hat', 5475), - ('Bench', 5310), - ('Bookcase', 5163), - ('Bee', 5140), - ('Computer monitor', 5078), - ('Hiking equipment', 4983), - ('Office building', 4981), - ('Coffee cup', 4748), - ('Curtain', 4685), - ('Plate', 4651), - ('Box', 4621), - ('Tomato', 4595), - ('Coffee table', 4529), - ('Office supplies', 4473), - ('Maple', 4416), - ('Muffin', 4365), - ('Cocktail', 4234), - ('Castle', 4197), - ('Couch', 4134), - ('Pumpkin', 3983), - ('Computer keyboard', 3960), - ('Human mouth', 3926), - ('Christmas tree', 3893), - ('Mushroom', 3883), - ('Swimming pool', 3809), - ('Pastry', 3799), - ('Lavender (Plant)', 3769), - ('Football helmet', 3732), - ('Bread', 3648), - ('Traffic sign', 3628), - ('Common sunflower', 3597), - ('Television', 3550), - ('Bed', 3525), - ('Cookie', 3485), - ('Fountain', 3484), - ('Paddle', 3447), - ('Bicycle helmet', 3429), - ('Porch', 3420), - ('Deer', 3387), - ('Fedora', 3339), - ('Canoe', 3338), - ('Carnivore', 3266), - ('Bowl', 3202), - ('Human eye', 3166), - ('Ball', 3118), - ('Pillow', 3077), - ('Salad', 3061), - ('Beetle', 3060), - ('Orange', 3050), - ('Drawer', 2958), - ('Platter', 2937), - ('Elephant', 2921), - ('Seafood', 2921), - ('Monkey', 2915), - ('Countertop', 2879), - ('Watercraft', 2831), - ('Helicopter', 2805), - ('Kitchen appliance', 2797), - ('Personal flotation device', 2781), - ('Swan', 2739), - ('Lamp', 2711), - ('Boot', 2695), - ('Bronze sculpture', 2693), - ('Chicken', 2677), - ('Taxi', 2643), - ('Juice', 2615), - ('Cowboy hat', 2604), - ('Apple', 2600), - ('Tin can', 2590), - ('Necklace', 2564), - ('Ice cream', 2560), - ('Human beard', 2539), - ('Coin', 2536), - ('Candle', 2515), - ('Cart', 2512), - ('High heels', 2441), - ('Weapon', 2433), - ('Handbag', 2406), - ('Penguin', 2396), - ('Rifle', 2352), - ('Violin', 2336), - ('Skull', 2304), - ('Lantern', 2285), - ('Scarf', 2269), - ('Saucer', 2225), - ('Sheep', 2215), - ('Vase', 2189), - ('Lily', 2180), - ('Mug', 2154), - ('Parrot', 2140), - ('Human ear', 2137), - ('Sandal', 2115), - ('Lizard', 2100), - ('Kitchen & dining room table', 2063), - ('Spider', 1977), - ('Coffee', 1974), - ('Goat', 1926), - ('Squirrel', 1922), - ('Cello', 1913), - ('Sushi', 1881), - ('Tortoise', 1876), - ('Pizza', 1870), - ('Studio couch', 1864), - ('Barrel', 1862), - ('Cosmetics', 1841), - ('Moths and butterflies', 1841), - ('Convenience store', 1817), - ('Watch', 1792), - ('Home appliance', 1786), - ('Harbor seal', 1780), - ('Luggage and bags', 1756), - ('Vehicle registration plate', 1754), - ('Shrimp', 1751), - ('Jellyfish', 1730), - ('French fries', 1723), - ('Egg (Food)', 1698), - ('Football', 1697), - ('Musical keyboard', 1683), - ('Falcon', 1674), - ('Candy', 1660), - ('Medical equipment', 1654), - ('Eagle', 1651), - ('Dinosaur', 1634), - ('Surfboard', 1630), - ('Tank', 1628), - ('Grape', 1624), - ('Lion', 1624), - ('Owl', 1622), - ('Ski', 1613), - ('Waste container', 1606), - ('Frog', 1591), - ('Sparrow', 1585), - ('Rabbit', 1581), - ('Pen', 1546), - ('Sea lion', 1537), - ('Spoon', 1521), - ('Sink', 1512), - ('Teddy bear', 1507), - ('Bull', 1495), - ('Sofa bed', 1490), - ('Dragonfly', 1479), - ('Brassiere', 1478), - ('Chest of drawers', 1472), - ('Aircraft', 1466), - ('Human foot', 1463), - ('Pig', 1455), - ('Fork', 1454), - ('Antelope', 1438), - ('Tripod', 1427), - ('Tool', 1424), - ('Cheese', 1422), - ('Lemon', 1397), - ('Hamburger', 1393), - ('Dolphin', 1390), - ('Mirror', 1390), - ('Marine mammal', 1387), - ('Giraffe', 1385), - ('Snake', 1368), - ('Gondola', 1364), - ('Wheelchair', 1360), - ('Piano', 1358), - ('Cupboard', 1348), - ('Banana', 1345), - ('Trumpet', 1335), - ('Lighthouse', 1333), - ('Invertebrate', 1317), - ('Carrot', 1268), - ('Sock', 1260), - ('Tiger', 1241), - ('Camel', 1224), - ('Parachute', 1224), - ('Bathroom accessory', 1223), - ('Earrings', 1221), - ('Headphones', 1218), - ('Skirt', 1198), - ('Skateboard', 1190), - ('Sandwich', 1148), - ('Saxophone', 1141), - ('Goldfish', 1136), - ('Stool', 1104), - ('Traffic light', 1097), - ('Shellfish', 1081), - ('Backpack', 1079), - ('Sea turtle', 1078), - ('Cucumber', 1075), - ('Tea', 1051), - ('Toilet', 1047), - ('Roller skates', 1040), - ('Mule', 1039), - ('Bust', 1031), - ('Broccoli', 1030), - ('Crab', 1020), - ('Oyster', 1019), - ('Cannon', 1012), - ('Zebra', 1012), - ('French horn', 1008), - ('Grapefruit', 998), - ('Whiteboard', 997), - ('Zucchini', 997), - ('Crocodile', 992), - - ('Clock', 960), - ('Wall clock', 958), - - ('Doughnut', 869), - ('Snail', 868), - - ('Baseball glove', 859), - - ('Panda', 830), - ('Tennis racket', 830), - - ('Pear', 652), - - ('Bagel', 617), - ('Oven', 616), - ('Ladybug', 615), - ('Shark', 615), - ('Polar bear', 614), - ('Ostrich', 609), - - ('Hot dog', 473), - ('Microwave oven', 467), - ('Fire hydrant', 20), - ('Stop sign', 20), - ('Parking meter', 20), - ('Bear', 20), - ('Flying disc', 20), - ('Snowboard', 20), - ('Tennis ball', 20), - ('Kite', 20), - ('Baseball bat', 20), - ('Kitchen knife', 20), - ('Knife', 20), - ('Submarine sandwich', 20), - ('Computer mouse', 20), - ('Remote control', 20), - ('Toaster', 20), - ('Sink', 20), - ('Refrigerator', 20), - ('Alarm clock', 20), - ('Wall clock', 20), - ('Scissors', 20), - ('Hair dryer', 20), - ('Toothbrush', 20), - ('Suitcase', 20) -] diff --git a/taming/data/sflckr.py b/taming/data/sflckr.py deleted file mode 100644 index 91101be..0000000 --- a/taming/data/sflckr.py +++ /dev/null @@ -1,91 +0,0 @@ -import os -import numpy as np -import cv2 -import albumentations -from PIL import Image -from torch.utils.data import Dataset - - -class SegmentationBase(Dataset): - def __init__(self, - data_csv, data_root, segmentation_root, - size=None, random_crop=False, interpolation="bicubic", - n_labels=182, shift_segmentation=False, - ): - self.n_labels = n_labels - self.shift_segmentation = shift_segmentation - self.data_csv = data_csv - self.data_root = data_root - self.segmentation_root = segmentation_root - with open(self.data_csv, "r") as f: - self.image_paths = f.read().splitlines() - self._length = len(self.image_paths) - self.labels = { - "relative_file_path_": [l for l in self.image_paths], - "file_path_": [os.path.join(self.data_root, l) - for l in self.image_paths], - "segmentation_path_": [os.path.join(self.segmentation_root, l.replace(".jpg", ".png")) - for l in self.image_paths] - } - - size = None if size is not None and size<=0 else size - self.size = size - if self.size is not None: - self.interpolation = interpolation - self.interpolation = { - "nearest": cv2.INTER_NEAREST, - "bilinear": cv2.INTER_LINEAR, - "bicubic": cv2.INTER_CUBIC, - "area": cv2.INTER_AREA, - "lanczos": cv2.INTER_LANCZOS4}[self.interpolation] - self.image_rescaler = albumentations.SmallestMaxSize(max_size=self.size, - interpolation=self.interpolation) - self.segmentation_rescaler = albumentations.SmallestMaxSize(max_size=self.size, - interpolation=cv2.INTER_NEAREST) - self.center_crop = not random_crop - if self.center_crop: - self.cropper = albumentations.CenterCrop(height=self.size, width=self.size) - else: - self.cropper = albumentations.RandomCrop(height=self.size, width=self.size) - self.preprocessor = self.cropper - - def __len__(self): - return self._length - - def __getitem__(self, i): - example = dict((k, self.labels[k][i]) for k in self.labels) - image = Image.open(example["file_path_"]) - if not image.mode == "RGB": - image = image.convert("RGB") - image = np.array(image).astype(np.uint8) - if self.size is not None: - image = self.image_rescaler(image=image)["image"] - segmentation = Image.open(example["segmentation_path_"]) - assert segmentation.mode == "L", segmentation.mode - segmentation = np.array(segmentation).astype(np.uint8) - if self.shift_segmentation: - # used to support segmentations containing unlabeled==255 label - segmentation = segmentation+1 - if self.size is not None: - segmentation = self.segmentation_rescaler(image=segmentation)["image"] - if self.size is not None: - processed = self.preprocessor(image=image, - mask=segmentation - ) - else: - processed = {"image": image, - "mask": segmentation - } - example["image"] = (processed["image"]/127.5 - 1.0).astype(np.float32) - segmentation = processed["mask"] - onehot = np.eye(self.n_labels)[segmentation] - example["segmentation"] = onehot - return example - - -class Examples(SegmentationBase): - def __init__(self, size=None, random_crop=False, interpolation="bicubic"): - super().__init__(data_csv="data/sflckr_examples.txt", - data_root="data/sflckr_images", - segmentation_root="data/sflckr_segmentations", - size=size, random_crop=random_crop, interpolation=interpolation) diff --git a/taming/data/utils.py b/taming/data/utils.py deleted file mode 100644 index 2b3c3d5..0000000 --- a/taming/data/utils.py +++ /dev/null @@ -1,169 +0,0 @@ -import collections -import os -import tarfile -import urllib -import zipfile -from pathlib import Path - -import numpy as np -import torch -from taming.data.helper_types import Annotation -from torch._six import string_classes -from torch.utils.data._utils.collate import np_str_obj_array_pattern, default_collate_err_msg_format -from tqdm import tqdm - - -def unpack(path): - if path.endswith("tar.gz"): - with tarfile.open(path, "r:gz") as tar: - tar.extractall(path=os.path.split(path)[0]) - elif path.endswith("tar"): - with tarfile.open(path, "r:") as tar: - tar.extractall(path=os.path.split(path)[0]) - elif path.endswith("zip"): - with zipfile.ZipFile(path, "r") as f: - f.extractall(path=os.path.split(path)[0]) - else: - raise NotImplementedError( - "Unknown file extension: {}".format(os.path.splitext(path)[1]) - ) - - -def reporthook(bar): - """tqdm progress bar for downloads.""" - - def hook(b=1, bsize=1, tsize=None): - if tsize is not None: - bar.total = tsize - bar.update(b * bsize - bar.n) - - return hook - - -def get_root(name): - base = "data/" - root = os.path.join(base, name) - os.makedirs(root, exist_ok=True) - return root - - -def is_prepared(root): - return Path(root).joinpath(".ready").exists() - - -def mark_prepared(root): - Path(root).joinpath(".ready").touch() - - -def prompt_download(file_, source, target_dir, content_dir=None): - targetpath = os.path.join(target_dir, file_) - while not os.path.exists(targetpath): - if content_dir is not None and os.path.exists( - os.path.join(target_dir, content_dir) - ): - break - print( - "Please download '{}' from '{}' to '{}'.".format(file_, source, targetpath) - ) - if content_dir is not None: - print( - "Or place its content into '{}'.".format( - os.path.join(target_dir, content_dir) - ) - ) - input("Press Enter when done...") - return targetpath - - -def download_url(file_, url, target_dir): - targetpath = os.path.join(target_dir, file_) - os.makedirs(target_dir, exist_ok=True) - with tqdm( - unit="B", unit_scale=True, unit_divisor=1024, miniters=1, desc=file_ - ) as bar: - urllib.request.urlretrieve(url, targetpath, reporthook=reporthook(bar)) - return targetpath - - -def download_urls(urls, target_dir): - paths = dict() - for fname, url in urls.items(): - outpath = download_url(fname, url, target_dir) - paths[fname] = outpath - return paths - - -def quadratic_crop(x, bbox, alpha=1.0): - """bbox is xmin, ymin, xmax, ymax""" - im_h, im_w = x.shape[:2] - bbox = np.array(bbox, dtype=np.float32) - bbox = np.clip(bbox, 0, max(im_h, im_w)) - center = 0.5 * (bbox[0] + bbox[2]), 0.5 * (bbox[1] + bbox[3]) - w = bbox[2] - bbox[0] - h = bbox[3] - bbox[1] - l = int(alpha * max(w, h)) - l = max(l, 2) - - required_padding = -1 * min( - center[0] - l, center[1] - l, im_w - (center[0] + l), im_h - (center[1] + l) - ) - required_padding = int(np.ceil(required_padding)) - if required_padding > 0: - padding = [ - [required_padding, required_padding], - [required_padding, required_padding], - ] - padding += [[0, 0]] * (len(x.shape) - 2) - x = np.pad(x, padding, "reflect") - center = center[0] + required_padding, center[1] + required_padding - xmin = int(center[0] - l / 2) - ymin = int(center[1] - l / 2) - return np.array(x[ymin : ymin + l, xmin : xmin + l, ...]) - - -def custom_collate(batch): - r"""source: pytorch 1.9.0, only one modification to original code """ - - elem = batch[0] - elem_type = type(elem) - if isinstance(elem, torch.Tensor): - out = None - if torch.utils.data.get_worker_info() is not None: - # If we're in a background process, concatenate directly into a - # shared memory tensor to avoid an extra copy - numel = sum([x.numel() for x in batch]) - storage = elem.storage()._new_shared(numel) - out = elem.new(storage) - return torch.stack(batch, 0, out=out) - elif elem_type.__module__ == 'numpy' and elem_type.__name__ != 'str_' \ - and elem_type.__name__ != 'string_': - if elem_type.__name__ == 'ndarray' or elem_type.__name__ == 'memmap': - # array of string classes and object - if np_str_obj_array_pattern.search(elem.dtype.str) is not None: - raise TypeError(default_collate_err_msg_format.format(elem.dtype)) - - return custom_collate([torch.as_tensor(b) for b in batch]) - elif elem.shape == (): # scalars - return torch.as_tensor(batch) - elif isinstance(elem, float): - return torch.tensor(batch, dtype=torch.float64) - elif isinstance(elem, int): - return torch.tensor(batch) - elif isinstance(elem, string_classes): - return batch - elif isinstance(elem, collections.abc.Mapping): - return {key: custom_collate([d[key] for d in batch]) for key in elem} - elif isinstance(elem, tuple) and hasattr(elem, '_fields'): # namedtuple - return elem_type(*(custom_collate(samples) for samples in zip(*batch))) - if isinstance(elem, collections.abc.Sequence) and isinstance(elem[0], Annotation): # added - return batch # added - elif isinstance(elem, collections.abc.Sequence): - # check to make sure that the elements in batch have consistent size - it = iter(batch) - elem_size = len(next(it)) - if not all(len(elem) == elem_size for elem in it): - raise RuntimeError('each element in list of batch should be of equal size') - transposed = zip(*batch) - return [custom_collate(samples) for samples in transposed] - - raise TypeError(default_collate_err_msg_format.format(elem_type)) diff --git a/taming/lr_scheduler.py b/taming/lr_scheduler.py deleted file mode 100644 index e598ed1..0000000 --- a/taming/lr_scheduler.py +++ /dev/null @@ -1,34 +0,0 @@ -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): - 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): - return self.schedule(n) - diff --git a/taming/models/cond_transformer.py b/taming/models/cond_transformer.py deleted file mode 100644 index e4c6373..0000000 --- a/taming/models/cond_transformer.py +++ /dev/null @@ -1,352 +0,0 @@ -import os, math -import torch -import torch.nn.functional as F -import pytorch_lightning as pl - -from main import instantiate_from_config -from taming.modules.util import SOSProvider - - -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 Net2NetTransformer(pl.LightningModule): - def __init__(self, - transformer_config, - first_stage_config, - cond_stage_config, - permuter_config=None, - ckpt_path=None, - ignore_keys=[], - first_stage_key="image", - cond_stage_key="depth", - downsample_cond_size=-1, - pkeep=1.0, - sos_token=0, - unconditional=False, - ): - super().__init__() - self.be_unconditional = unconditional - self.sos_token = sos_token - self.first_stage_key = first_stage_key - self.cond_stage_key = cond_stage_key - self.init_first_stage_from_ckpt(first_stage_config) - self.init_cond_stage_from_ckpt(cond_stage_config) - if permuter_config is None: - permuter_config = {"target": "taming.modules.transformer.permuter.Identity"} - self.permuter = instantiate_from_config(config=permuter_config) - self.transformer = instantiate_from_config(config=transformer_config) - - if ckpt_path is not None: - self.init_from_ckpt(ckpt_path, ignore_keys=ignore_keys) - self.downsample_cond_size = downsample_cond_size - self.pkeep = pkeep - - def init_from_ckpt(self, path, ignore_keys=list()): - sd = torch.load(path, map_location="cpu")["state_dict"] - for k in sd.keys(): - for ik in ignore_keys: - if k.startswith(ik): - self.print("Deleting key {} from state_dict.".format(k)) - del sd[k] - self.load_state_dict(sd, strict=False) - print(f"Restored from {path}") - - def init_first_stage_from_ckpt(self, config): - model = instantiate_from_config(config) - model = model.eval() - model.train = disabled_train - self.first_stage_model = model - - def init_cond_stage_from_ckpt(self, config): - if config == "__is_first_stage__": - print("Using first stage also as cond stage.") - self.cond_stage_model = self.first_stage_model - elif config == "__is_unconditional__" or self.be_unconditional: - print(f"Using no cond stage. Assuming the training is intended to be unconditional. " - f"Prepending {self.sos_token} as a sos token.") - self.be_unconditional = True - self.cond_stage_key = self.first_stage_key - self.cond_stage_model = SOSProvider(self.sos_token) - else: - model = instantiate_from_config(config) - model = model.eval() - model.train = disabled_train - self.cond_stage_model = model - - def forward(self, x, c): - # one step to produce the logits - _, z_indices = self.encode_to_z(x) - _, c_indices = self.encode_to_c(c) - - if self.training and self.pkeep < 1.0: - mask = torch.bernoulli(self.pkeep*torch.ones(z_indices.shape, - device=z_indices.device)) - mask = mask.round().to(dtype=torch.int64) - r_indices = torch.randint_like(z_indices, self.transformer.config.vocab_size) - a_indices = mask*z_indices+(1-mask)*r_indices - else: - a_indices = z_indices - - cz_indices = torch.cat((c_indices, a_indices), dim=1) - - # target includes all sequence elements (no need to handle first one - # differently because we are conditioning) - target = z_indices - # make the prediction - logits, _ = self.transformer(cz_indices[:, :-1]) - # cut off conditioning outputs - output i corresponds to p(z_i | z_{ -1: - c = F.interpolate(c, size=(self.downsample_cond_size, self.downsample_cond_size)) - quant_c, _, [_,_,indices] = self.cond_stage_model.encode(c) - if len(indices.shape) > 2: - indices = indices.view(c.shape[0], -1) - return quant_c, indices - - @torch.no_grad() - def decode_to_img(self, index, zshape): - index = self.permuter(index, reverse=True) - bhwc = (zshape[0],zshape[2],zshape[3],zshape[1]) - quant_z = self.first_stage_model.quantize.get_codebook_entry( - index.reshape(-1), shape=bhwc) - x = self.first_stage_model.decode(quant_z) - return x - - @torch.no_grad() - def log_images(self, batch, temperature=None, top_k=None, callback=None, lr_interface=False, **kwargs): - log = dict() - - N = 4 - if lr_interface: - x, c = self.get_xc(batch, N, diffuse=False, upsample_factor=8) - else: - x, c = self.get_xc(batch, N) - x = x.to(device=self.device) - c = c.to(device=self.device) - - quant_z, z_indices = self.encode_to_z(x) - quant_c, c_indices = self.encode_to_c(c) - - # create a "half"" sample - z_start_indices = z_indices[:,:z_indices.shape[1]//2] - index_sample = self.sample(z_start_indices, c_indices, - steps=z_indices.shape[1]-z_start_indices.shape[1], - temperature=temperature if temperature is not None else 1.0, - sample=True, - top_k=top_k if top_k is not None else 100, - callback=callback if callback is not None else lambda k: None) - x_sample = self.decode_to_img(index_sample, quant_z.shape) - - # sample - z_start_indices = z_indices[:, :0] - index_sample = self.sample(z_start_indices, c_indices, - steps=z_indices.shape[1], - temperature=temperature if temperature is not None else 1.0, - sample=True, - top_k=top_k if top_k is not None else 100, - callback=callback if callback is not None else lambda k: None) - x_sample_nopix = self.decode_to_img(index_sample, quant_z.shape) - - # det sample - z_start_indices = z_indices[:, :0] - index_sample = self.sample(z_start_indices, c_indices, - steps=z_indices.shape[1], - sample=False, - callback=callback if callback is not None else lambda k: None) - x_sample_det = self.decode_to_img(index_sample, quant_z.shape) - - # reconstruction - x_rec = self.decode_to_img(z_indices, quant_z.shape) - - log["inputs"] = x - log["reconstructions"] = x_rec - - if self.cond_stage_key in ["objects_bbox", "objects_center_points"]: - figure_size = (x_rec.shape[2], x_rec.shape[3]) - dataset = kwargs["pl_module"].trainer.datamodule.datasets["validation"] - label_for_category_no = dataset.get_textual_label_for_category_no - plotter = dataset.conditional_builders[self.cond_stage_key].plot - log["conditioning"] = torch.zeros_like(log["reconstructions"]) - for i in range(quant_c.shape[0]): - log["conditioning"][i] = plotter(quant_c[i], label_for_category_no, figure_size) - log["conditioning_rec"] = log["conditioning"] - elif self.cond_stage_key != "image": - cond_rec = self.cond_stage_model.decode(quant_c) - if self.cond_stage_key == "segmentation": - # get image from segmentation mask - num_classes = cond_rec.shape[1] - - c = torch.argmax(c, dim=1, keepdim=True) - c = F.one_hot(c, num_classes=num_classes) - c = c.squeeze(1).permute(0, 3, 1, 2).float() - c = self.cond_stage_model.to_rgb(c) - - cond_rec = torch.argmax(cond_rec, dim=1, keepdim=True) - cond_rec = F.one_hot(cond_rec, num_classes=num_classes) - cond_rec = cond_rec.squeeze(1).permute(0, 3, 1, 2).float() - cond_rec = self.cond_stage_model.to_rgb(cond_rec) - log["conditioning_rec"] = cond_rec - log["conditioning"] = c - - log["samples_half"] = x_sample - log["samples_nopix"] = x_sample_nopix - log["samples_det"] = x_sample_det - return log - - def get_input(self, key, batch): - x = batch[key] - if len(x.shape) == 3: - x = x[..., None] - if len(x.shape) == 4: - x = x.permute(0, 3, 1, 2).to(memory_format=torch.contiguous_format) - if x.dtype == torch.double: - x = x.float() - return x - - def get_xc(self, batch, N=None): - x = self.get_input(self.first_stage_key, batch) - c = self.get_input(self.cond_stage_key, batch) - if N is not None: - x = x[:N] - c = c[:N] - return x, c - - def shared_step(self, batch, batch_idx): - x, c = self.get_xc(batch) - logits, target = self(x, c) - loss = F.cross_entropy(logits.reshape(-1, logits.size(-1)), target.reshape(-1)) - return loss - - def training_step(self, batch, batch_idx): - loss = self.shared_step(batch, batch_idx) - self.log("train/loss", loss, prog_bar=True, logger=True, on_step=True, on_epoch=True) - return loss - - def validation_step(self, batch, batch_idx): - loss = self.shared_step(batch, batch_idx) - self.log("val/loss", loss, prog_bar=True, logger=True, on_step=True, on_epoch=True) - return loss - - def configure_optimizers(self): - """ - Following minGPT: - This long function is unfortunately doing something very simple and is being very defensive: - We are separating out all parameters of the model into two buckets: those that will experience - weight decay for regularization and those that won't (biases, and layernorm/embedding weights). - We are then returning the PyTorch optimizer object. - """ - # separate out all parameters to those that will and won't experience regularizing weight decay - decay = set() - no_decay = set() - whitelist_weight_modules = (torch.nn.Linear, ) - blacklist_weight_modules = (torch.nn.LayerNorm, torch.nn.Embedding) - for mn, m in self.transformer.named_modules(): - for pn, p in m.named_parameters(): - fpn = '%s.%s' % (mn, pn) if mn else pn # full param name - - if pn.endswith('bias'): - # all biases will not be decayed - no_decay.add(fpn) - elif pn.endswith('weight') and isinstance(m, whitelist_weight_modules): - # weights of whitelist modules will be weight decayed - decay.add(fpn) - elif pn.endswith('weight') and isinstance(m, blacklist_weight_modules): - # weights of blacklist modules will NOT be weight decayed - no_decay.add(fpn) - - # special case the position embedding parameter in the root GPT module as not decayed - no_decay.add('pos_emb') - - # validate that we considered every parameter - param_dict = {pn: p for pn, p in self.transformer.named_parameters()} - inter_params = decay & no_decay - union_params = decay | no_decay - assert len(inter_params) == 0, "parameters %s made it into both decay/no_decay sets!" % (str(inter_params), ) - assert len(param_dict.keys() - union_params) == 0, "parameters %s were not separated into either decay/no_decay set!" \ - % (str(param_dict.keys() - union_params), ) - - # create the pytorch optimizer object - optim_groups = [ - {"params": [param_dict[pn] for pn in sorted(list(decay))], "weight_decay": 0.01}, - {"params": [param_dict[pn] for pn in sorted(list(no_decay))], "weight_decay": 0.0}, - ] - optimizer = torch.optim.AdamW(optim_groups, lr=self.learning_rate, betas=(0.9, 0.95)) - return optimizer diff --git a/taming/models/dummy_cond_stage.py b/taming/models/dummy_cond_stage.py deleted file mode 100644 index 6e19938..0000000 --- a/taming/models/dummy_cond_stage.py +++ /dev/null @@ -1,22 +0,0 @@ -from torch import Tensor - - -class DummyCondStage: - def __init__(self, conditional_key): - self.conditional_key = conditional_key - self.train = None - - def eval(self): - return self - - @staticmethod - def encode(c: Tensor): - return c, None, (None, None, c) - - @staticmethod - def decode(c: Tensor): - return c - - @staticmethod - def to_rgb(c: Tensor): - return c diff --git a/taming/models/vqgan.py b/taming/models/vqgan.py deleted file mode 100644 index a6950ba..0000000 --- a/taming/models/vqgan.py +++ /dev/null @@ -1,404 +0,0 @@ -import torch -import torch.nn.functional as F -import pytorch_lightning as pl - -from main import instantiate_from_config - -from taming.modules.diffusionmodules.model import Encoder, Decoder -from taming.modules.vqvae.quantize import VectorQuantizer2 as VectorQuantizer -from taming.modules.vqvae.quantize import GumbelQuantize -from taming.modules.vqvae.quantize import EMAVectorQuantizer - -class VQModel(pl.LightningModule): - def __init__(self, - ddconfig, - lossconfig, - n_embed, - embed_dim, - ckpt_path=None, - ignore_keys=[], - image_key="image", - colorize_nlabels=None, - monitor=None, - remap=None, - sane_index_shape=False, # tell vector quantizer to return indices as bhw - ): - super().__init__() - self.image_key = image_key - self.encoder = Encoder(**ddconfig) - self.decoder = Decoder(**ddconfig) - 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 ckpt_path is not None: - self.init_from_ckpt(ckpt_path, ignore_keys=ignore_keys) - self.image_key = image_key - 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 - - 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) - h = self.quant_conv(h) - quant, emb_loss, info = self.quantize(h) - return quant, emb_loss, info - - 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): - quant, diff, _ = self.encode(input) - dec = self.decode(quant) - return dec, diff - - def get_input(self, batch, k): - x = batch[k] - if len(x.shape) == 3: - x = x[..., None] - x = x.permute(0, 3, 1, 2).to(memory_format=torch.contiguous_format) - return x.float() - - def training_step(self, batch, batch_idx, optimizer_idx): - x = self.get_input(batch, self.image_key) - xrec, qloss = self(x) - - if optimizer_idx == 0: - # autoencode - aeloss, log_dict_ae = self.loss(qloss, x, xrec, optimizer_idx, self.global_step, - last_layer=self.get_last_layer(), split="train") - - self.log("train/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=True) - return aeloss - - if optimizer_idx == 1: - # discriminator - discloss, log_dict_disc = self.loss(qloss, x, xrec, optimizer_idx, self.global_step, - last_layer=self.get_last_layer(), split="train") - self.log("train/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=True) - return discloss - - def validation_step(self, batch, batch_idx): - x = self.get_input(batch, self.image_key) - xrec, qloss = self(x) - aeloss, log_dict_ae = self.loss(qloss, x, xrec, 0, self.global_step, - last_layer=self.get_last_layer(), split="val") - - discloss, log_dict_disc = self.loss(qloss, x, xrec, 1, self.global_step, - last_layer=self.get_last_layer(), split="val") - rec_loss = log_dict_ae["val/rec_loss"] - self.log("val/rec_loss", rec_loss, - prog_bar=True, logger=True, on_step=True, on_epoch=True, sync_dist=True) - self.log("val/aeloss", aeloss, - prog_bar=True, logger=True, on_step=True, on_epoch=True, sync_dist=True) - 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.quantize.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 - - def log_images(self, batch, **kwargs): - log = dict() - x = self.get_input(batch, self.image_key) - x = x.to(self.device) - xrec, _ = 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["inputs"] = x - log["reconstructions"] = xrec - 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 VQSegmentationModel(VQModel): - def __init__(self, n_labels, *args, **kwargs): - super().__init__(*args, **kwargs) - self.register_buffer("colorize", torch.randn(3, n_labels, 1, 1)) - - def configure_optimizers(self): - lr = self.learning_rate - 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, betas=(0.5, 0.9)) - return opt_ae - - def training_step(self, batch, batch_idx): - x = self.get_input(batch, self.image_key) - xrec, qloss = self(x) - aeloss, log_dict_ae = self.loss(qloss, x, xrec, split="train") - self.log_dict(log_dict_ae, prog_bar=False, logger=True, on_step=True, on_epoch=True) - return aeloss - - def validation_step(self, batch, batch_idx): - x = self.get_input(batch, self.image_key) - xrec, qloss = self(x) - aeloss, log_dict_ae = self.loss(qloss, x, xrec, split="val") - self.log_dict(log_dict_ae, prog_bar=False, logger=True, on_step=True, on_epoch=True) - total_loss = log_dict_ae["val/total_loss"] - self.log("val/total_loss", total_loss, - prog_bar=True, logger=True, on_step=True, on_epoch=True, sync_dist=True) - return aeloss - - @torch.no_grad() - def log_images(self, batch, **kwargs): - log = dict() - x = self.get_input(batch, self.image_key) - x = x.to(self.device) - xrec, _ = self(x) - if x.shape[1] > 3: - # colorize with random projection - assert xrec.shape[1] > 3 - # convert logits to indices - xrec = torch.argmax(xrec, dim=1, keepdim=True) - xrec = F.one_hot(xrec, num_classes=x.shape[1]) - xrec = xrec.squeeze(1).permute(0, 3, 1, 2).float() - x = self.to_rgb(x) - xrec = self.to_rgb(xrec) - log["inputs"] = x - log["reconstructions"] = xrec - return log - - -class VQNoDiscModel(VQModel): - def __init__(self, - ddconfig, - lossconfig, - n_embed, - embed_dim, - ckpt_path=None, - ignore_keys=[], - image_key="image", - colorize_nlabels=None - ): - super().__init__(ddconfig=ddconfig, lossconfig=lossconfig, n_embed=n_embed, embed_dim=embed_dim, - ckpt_path=ckpt_path, ignore_keys=ignore_keys, image_key=image_key, - colorize_nlabels=colorize_nlabels) - - def training_step(self, batch, batch_idx): - x = self.get_input(batch, self.image_key) - xrec, qloss = self(x) - # autoencode - aeloss, log_dict_ae = self.loss(qloss, x, xrec, self.global_step, split="train") - output = pl.TrainResult(minimize=aeloss) - output.log("train/aeloss", aeloss, - prog_bar=True, logger=True, on_step=True, on_epoch=True) - output.log_dict(log_dict_ae, prog_bar=False, logger=True, on_step=True, on_epoch=True) - return output - - def validation_step(self, batch, batch_idx): - x = self.get_input(batch, self.image_key) - xrec, qloss = self(x) - aeloss, log_dict_ae = self.loss(qloss, x, xrec, self.global_step, split="val") - rec_loss = log_dict_ae["val/rec_loss"] - output = pl.EvalResult(checkpoint_on=rec_loss) - output.log("val/rec_loss", rec_loss, - prog_bar=True, logger=True, on_step=True, on_epoch=True) - output.log("val/aeloss", aeloss, - prog_bar=True, logger=True, on_step=True, on_epoch=True) - output.log_dict(log_dict_ae) - - return output - - def configure_optimizers(self): - optimizer = 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=self.learning_rate, betas=(0.5, 0.9)) - return optimizer - - -class GumbelVQ(VQModel): - def __init__(self, - ddconfig, - lossconfig, - n_embed, - embed_dim, - temperature_scheduler_config, - ckpt_path=None, - ignore_keys=[], - image_key="image", - colorize_nlabels=None, - monitor=None, - kl_weight=1e-8, - remap=None, - ): - - z_channels = ddconfig["z_channels"] - super().__init__(ddconfig, - lossconfig, - n_embed, - embed_dim, - ckpt_path=None, - ignore_keys=ignore_keys, - image_key=image_key, - colorize_nlabels=colorize_nlabels, - monitor=monitor, - ) - - self.loss.n_classes = n_embed - self.vocab_size = n_embed - - self.quantize = GumbelQuantize(z_channels, embed_dim, - n_embed=n_embed, - kl_weight=kl_weight, temp_init=1.0, - remap=remap) - - self.temperature_scheduler = instantiate_from_config(temperature_scheduler_config) # annealing of temp - - if ckpt_path is not None: - self.init_from_ckpt(ckpt_path, ignore_keys=ignore_keys) - - def temperature_scheduling(self): - self.quantize.temperature = self.temperature_scheduler(self.global_step) - - def encode_to_prequant(self, x): - h = self.encoder(x) - h = self.quant_conv(h) - return h - - def decode_code(self, code_b): - raise NotImplementedError - - def training_step(self, batch, batch_idx, optimizer_idx): - self.temperature_scheduling() - x = self.get_input(batch, self.image_key) - xrec, qloss = self(x) - - if optimizer_idx == 0: - # autoencode - aeloss, log_dict_ae = self.loss(qloss, x, xrec, optimizer_idx, self.global_step, - last_layer=self.get_last_layer(), split="train") - - self.log_dict(log_dict_ae, prog_bar=False, logger=True, on_step=True, on_epoch=True) - self.log("temperature", self.quantize.temperature, 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, xrec, optimizer_idx, self.global_step, - last_layer=self.get_last_layer(), split="train") - 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): - x = self.get_input(batch, self.image_key) - xrec, qloss = 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") - - discloss, log_dict_disc = self.loss(qloss, x, xrec, 1, self.global_step, - last_layer=self.get_last_layer(), split="val") - rec_loss = log_dict_ae["val/rec_loss"] - self.log("val/rec_loss", rec_loss, - prog_bar=True, logger=True, on_step=False, on_epoch=True, sync_dist=True) - self.log("val/aeloss", aeloss, - prog_bar=True, logger=True, on_step=False, on_epoch=True, sync_dist=True) - self.log_dict(log_dict_ae) - self.log_dict(log_dict_disc) - return self.log_dict - - def log_images(self, batch, **kwargs): - log = dict() - x = self.get_input(batch, self.image_key) - x = x.to(self.device) - # encode - h = self.encoder(x) - h = self.quant_conv(h) - quant, _, _ = self.quantize(h) - # decode - x_rec = self.decode(quant) - log["inputs"] = x - log["reconstructions"] = x_rec - return log - - -class EMAVQ(VQModel): - def __init__(self, - ddconfig, - lossconfig, - n_embed, - embed_dim, - ckpt_path=None, - ignore_keys=[], - image_key="image", - colorize_nlabels=None, - monitor=None, - remap=None, - sane_index_shape=False, # tell vector quantizer to return indices as bhw - ): - super().__init__(ddconfig, - lossconfig, - n_embed, - embed_dim, - ckpt_path=None, - ignore_keys=ignore_keys, - image_key=image_key, - colorize_nlabels=colorize_nlabels, - monitor=monitor, - ) - self.quantize = EMAVectorQuantizer(n_embed=n_embed, - embedding_dim=embed_dim, - beta=0.25, - remap=remap) - def configure_optimizers(self): - lr = self.learning_rate - #Remove self.quantize from parameter list since it is updated via EMA - 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], [] \ No newline at end of file diff --git a/taming/modules/diffusionmodules/model.py b/taming/modules/diffusionmodules/model.py deleted file mode 100644 index d3a5db6..0000000 --- a/taming/modules/diffusionmodules/model.py +++ /dev/null @@ -1,776 +0,0 @@ -# pytorch_diffusion + derived encoder decoder -import math -import torch -import torch.nn as nn -import numpy as np - - -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): - return torch.nn.GroupNorm(num_groups=32, 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 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_ - - -class Model(nn.Module): - def __init__(self, *, ch, out_ch, ch_mult=(1,2,4,8), num_res_blocks, - attn_resolutions, dropout=0.0, resamp_with_conv=True, in_channels, - resolution, use_timestep=True): - super().__init__() - self.ch = ch - self.temb_ch = self.ch*4 - self.num_resolutions = len(ch_mult) - self.num_res_blocks = num_res_blocks - self.resolution = resolution - 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) - - curr_res = resolution - 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 curr_res in attn_resolutions: - attn.append(AttnBlock(block_in)) - down = nn.Module() - down.block = block - down.attn = attn - if i_level != self.num_resolutions-1: - down.downsample = Downsample(block_in, resamp_with_conv) - curr_res = curr_res // 2 - 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 = AttnBlock(block_in) - 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 curr_res in attn_resolutions: - attn.append(AttnBlock(block_in)) - up = nn.Module() - up.block = block - up.attn = attn - if i_level != 0: - up.upsample = Upsample(block_in, resamp_with_conv) - curr_res = curr_res * 2 - 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): - #assert x.shape[2] == x.shape[3] == self.resolution - - 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 - - -class Encoder(nn.Module): - def __init__(self, *, ch, out_ch, ch_mult=(1,2,4,8), num_res_blocks, - attn_resolutions, dropout=0.0, resamp_with_conv=True, in_channels, - resolution, z_channels, double_z=True, **ignore_kwargs): - super().__init__() - self.ch = ch - self.temb_ch = 0 - self.num_resolutions = len(ch_mult) - self.num_res_blocks = num_res_blocks - self.resolution = resolution - self.in_channels = in_channels - - # downsampling - self.conv_in = torch.nn.Conv2d(in_channels, - self.ch, - kernel_size=3, - stride=1, - padding=1) - - curr_res = resolution - 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 curr_res in attn_resolutions: - attn.append(AttnBlock(block_in)) - down = nn.Module() - down.block = block - down.attn = attn - if i_level != self.num_resolutions-1: - down.downsample = Downsample(block_in, resamp_with_conv) - curr_res = curr_res // 2 - 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 = AttnBlock(block_in) - 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): - #assert x.shape[2] == x.shape[3] == self.resolution, "{}, {}, {}".format(x.shape[2], x.shape[3], self.resolution) - - # 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_resolutions, dropout=0.0, resamp_with_conv=True, in_channels, - resolution, z_channels, give_pre_end=False, **ignorekwargs): - super().__init__() - self.ch = ch - self.temb_ch = 0 - self.num_resolutions = len(ch_mult) - self.num_res_blocks = num_res_blocks - self.resolution = resolution - self.in_channels = in_channels - self.give_pre_end = give_pre_end - - # 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] - curr_res = resolution // 2**(self.num_resolutions-1) - self.z_shape = (1,z_channels,curr_res,curr_res) - print("Working with z of shape {} = {} dimensions.".format( - self.z_shape, np.prod(self.z_shape))) - - # 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 = AttnBlock(block_in) - 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 curr_res in attn_resolutions: - attn.append(AttnBlock(block_in)) - up = nn.Module() - up.block = block - up.attn = attn - if i_level != 0: - up.upsample = Upsample(block_in, resamp_with_conv) - curr_res = curr_res * 2 - 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): - #assert z.shape[1:] == self.z_shape[1:] - 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) - return h - - -class VUNet(nn.Module): - def __init__(self, *, ch, out_ch, ch_mult=(1,2,4,8), num_res_blocks, - attn_resolutions, dropout=0.0, resamp_with_conv=True, - in_channels, c_channels, - resolution, z_channels, use_timestep=False, **ignore_kwargs): - super().__init__() - self.ch = ch - self.temb_ch = self.ch*4 - self.num_resolutions = len(ch_mult) - self.num_res_blocks = num_res_blocks - self.resolution = resolution - - 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(c_channels, - self.ch, - kernel_size=3, - stride=1, - padding=1) - - curr_res = resolution - 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 curr_res in attn_resolutions: - attn.append(AttnBlock(block_in)) - down = nn.Module() - down.block = block - down.attn = attn - if i_level != self.num_resolutions-1: - down.downsample = Downsample(block_in, resamp_with_conv) - curr_res = curr_res // 2 - self.down.append(down) - - self.z_in = torch.nn.Conv2d(z_channels, - block_in, - kernel_size=1, - stride=1, - padding=0) - # middle - self.mid = nn.Module() - self.mid.block_1 = ResnetBlock(in_channels=2*block_in, - out_channels=block_in, - temb_channels=self.temb_ch, - dropout=dropout) - self.mid.attn_1 = AttnBlock(block_in) - 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 curr_res in attn_resolutions: - attn.append(AttnBlock(block_in)) - up = nn.Module() - up.block = block - up.attn = attn - if i_level != 0: - up.upsample = Upsample(block_in, resamp_with_conv) - curr_res = curr_res * 2 - 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, z): - #assert x.shape[2] == x.shape[3] == self.resolution - - 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] - z = self.z_in(z) - h = torch.cat((h,z),dim=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 - - -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, resolution, - 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 - curr_res = resolution // 2 ** (self.num_resolutions - 1) - 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)) - curr_res = curr_res * 2 - - # 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 - diff --git a/taming/modules/discriminator/model.py b/taming/modules/discriminator/model.py deleted file mode 100644 index 2aaa311..0000000 --- a/taming/modules/discriminator/model.py +++ /dev/null @@ -1,67 +0,0 @@ -import functools -import torch.nn as nn - - -from taming.modules.util import ActNorm - - -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) - - -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=3, 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 - 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 += [ - 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, 1, 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) diff --git a/taming/modules/losses/__init__.py b/taming/modules/losses/__init__.py deleted file mode 100644 index d09caf9..0000000 --- a/taming/modules/losses/__init__.py +++ /dev/null @@ -1,2 +0,0 @@ -from taming.modules.losses.vqperceptual import DummyLoss - diff --git a/taming/modules/losses/lpips.py b/taming/modules/losses/lpips.py deleted file mode 100644 index a728044..0000000 --- a/taming/modules/losses/lpips.py +++ /dev/null @@ -1,123 +0,0 @@ -"""Stripped version of https://github.com/richzhang/PerceptualSimilarity/tree/master/models""" - -import torch -import torch.nn as nn -from torchvision import models -from collections import namedtuple - -from taming.util import get_ckpt_path - - -class LPIPS(nn.Module): - # Learned perceptual metric - def __init__(self, use_dropout=True): - super().__init__() - self.scaling_layer = ScalingLayer() - self.chns = [64, 128, 256, 512, 512] # vg16 features - self.net = vgg16(pretrained=True, requires_grad=False) - self.lin0 = NetLinLayer(self.chns[0], use_dropout=use_dropout) - self.lin1 = NetLinLayer(self.chns[1], use_dropout=use_dropout) - self.lin2 = NetLinLayer(self.chns[2], use_dropout=use_dropout) - self.lin3 = NetLinLayer(self.chns[3], use_dropout=use_dropout) - self.lin4 = NetLinLayer(self.chns[4], use_dropout=use_dropout) - self.load_from_pretrained() - for param in self.parameters(): - param.requires_grad = False - - def load_from_pretrained(self, name="vgg_lpips"): - ckpt = get_ckpt_path(name, "taming/modules/autoencoder/lpips") - self.load_state_dict(torch.load(ckpt, map_location=torch.device("cpu")), strict=False) - print("loaded pretrained LPIPS loss from {}".format(ckpt)) - - @classmethod - def from_pretrained(cls, name="vgg_lpips"): - if name != "vgg_lpips": - raise NotImplementedError - model = cls() - ckpt = get_ckpt_path(name) - model.load_state_dict(torch.load(ckpt, map_location=torch.device("cpu")), strict=False) - return model - - def forward(self, input, target): - in0_input, in1_input = (self.scaling_layer(input), self.scaling_layer(target)) - outs0, outs1 = self.net(in0_input), self.net(in1_input) - feats0, feats1, diffs = {}, {}, {} - lins = [self.lin0, self.lin1, self.lin2, self.lin3, self.lin4] - for kk in range(len(self.chns)): - feats0[kk], feats1[kk] = normalize_tensor(outs0[kk]), normalize_tensor(outs1[kk]) - diffs[kk] = (feats0[kk] - feats1[kk]) ** 2 - - res = [spatial_average(lins[kk].model(diffs[kk]), keepdim=True) for kk in range(len(self.chns))] - val = res[0] - for l in range(1, len(self.chns)): - val += res[l] - return val - - -class ScalingLayer(nn.Module): - def __init__(self): - super(ScalingLayer, self).__init__() - self.register_buffer('shift', torch.Tensor([-.030, -.088, -.188])[None, :, None, None]) - self.register_buffer('scale', torch.Tensor([.458, .448, .450])[None, :, None, None]) - - def forward(self, inp): - return (inp - self.shift) / self.scale - - -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 vgg16(torch.nn.Module): - def __init__(self, requires_grad=False, pretrained=True): - super(vgg16, self).__init__() - vgg_pretrained_features = models.vgg16(pretrained=pretrained).features - self.slice1 = torch.nn.Sequential() - self.slice2 = torch.nn.Sequential() - self.slice3 = torch.nn.Sequential() - self.slice4 = torch.nn.Sequential() - self.slice5 = torch.nn.Sequential() - self.N_slices = 5 - for x in range(4): - self.slice1.add_module(str(x), vgg_pretrained_features[x]) - for x in range(4, 9): - self.slice2.add_module(str(x), vgg_pretrained_features[x]) - for x in range(9, 16): - self.slice3.add_module(str(x), vgg_pretrained_features[x]) - for x in range(16, 23): - self.slice4.add_module(str(x), vgg_pretrained_features[x]) - for x in range(23, 30): - self.slice5.add_module(str(x), vgg_pretrained_features[x]) - if not requires_grad: - for param in self.parameters(): - param.requires_grad = False - - def forward(self, X): - h = self.slice1(X) - h_relu1_2 = h - h = self.slice2(h) - h_relu2_2 = h - h = self.slice3(h) - h_relu3_3 = h - h = self.slice4(h) - h_relu4_3 = h - h = self.slice5(h) - h_relu5_3 = h - vgg_outputs = namedtuple("VggOutputs", ['relu1_2', 'relu2_2', 'relu3_3', 'relu4_3', 'relu5_3']) - out = vgg_outputs(h_relu1_2, h_relu2_2, h_relu3_3, h_relu4_3, h_relu5_3) - return out - - -def normalize_tensor(x,eps=1e-10): - norm_factor = torch.sqrt(torch.sum(x**2,dim=1,keepdim=True)) - return x/(norm_factor+eps) - - -def spatial_average(x, keepdim=True): - return x.mean([2,3],keepdim=keepdim) - diff --git a/taming/modules/losses/segmentation.py b/taming/modules/losses/segmentation.py deleted file mode 100644 index 4ba77de..0000000 --- a/taming/modules/losses/segmentation.py +++ /dev/null @@ -1,22 +0,0 @@ -import torch.nn as nn -import torch.nn.functional as F - - -class BCELoss(nn.Module): - def forward(self, prediction, target): - loss = F.binary_cross_entropy_with_logits(prediction,target) - return loss, {} - - -class BCELossWithQuant(nn.Module): - def __init__(self, codebook_weight=1.): - super().__init__() - self.codebook_weight = codebook_weight - - def forward(self, qloss, target, prediction, split): - bce_loss = F.binary_cross_entropy_with_logits(prediction,target) - loss = bce_loss + self.codebook_weight*qloss - return loss, {"{}/total_loss".format(split): loss.clone().detach().mean(), - "{}/bce_loss".format(split): bce_loss.detach().mean(), - "{}/quant_loss".format(split): qloss.detach().mean() - } diff --git a/taming/modules/losses/vqperceptual.py b/taming/modules/losses/vqperceptual.py deleted file mode 100644 index c2febd4..0000000 --- a/taming/modules/losses/vqperceptual.py +++ /dev/null @@ -1,136 +0,0 @@ -import torch -import torch.nn as nn -import torch.nn.functional as F - -from taming.modules.losses.lpips import LPIPS -from taming.modules.discriminator.model import NLayerDiscriminator, weights_init - - -class DummyLoss(nn.Module): - def __init__(self): - super().__init__() - - -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 - - -class VQLPIPSWithDiscriminator(nn.Module): - def __init__(self, disc_start, codebook_weight=1.0, pixelloss_weight=1.0, - disc_num_layers=3, disc_in_channels=3, disc_factor=1.0, disc_weight=1.0, - perceptual_weight=1.0, use_actnorm=False, disc_conditional=False, - disc_ndf=64, disc_loss="hinge"): - super().__init__() - assert disc_loss in ["hinge", "vanilla"] - self.codebook_weight = codebook_weight - self.pixel_weight = pixelloss_weight - self.perceptual_loss = LPIPS().eval() - self.perceptual_weight = perceptual_weight - - self.discriminator = NLayerDiscriminator(input_nc=disc_in_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"VQLPIPSWithDiscriminator running with {disc_loss} 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, codebook_loss, inputs, reconstructions, optimizer_idx, - global_step, last_layer=None, cond=None, split="train"): - 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 - else: - p_loss = torch.tensor([0.0]) - - nll_loss = rec_loss - #nll_loss = torch.sum(nll_loss) / nll_loss.shape[0] - nll_loss = torch.mean(nll_loss) - - # now the GAN part - 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) - - 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) - - disc_factor = adopt_weight(self.disc_factor, global_step, threshold=self.discriminator_iter_start) - 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(), - "{}/nll_loss".format(split): nll_loss.detach().mean(), - "{}/rec_loss".format(split): rec_loss.detach().mean(), - "{}/p_loss".format(split): p_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)) - - disc_factor = adopt_weight(self.disc_factor, global_step, threshold=self.discriminator_iter_start) - 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/taming/modules/misc/coord.py b/taming/modules/misc/coord.py deleted file mode 100644 index ee69b0c..0000000 --- a/taming/modules/misc/coord.py +++ /dev/null @@ -1,31 +0,0 @@ -import torch - -class CoordStage(object): - def __init__(self, n_embed, down_factor): - self.n_embed = n_embed - self.down_factor = down_factor - - def eval(self): - return self - - def encode(self, c): - """fake vqmodel interface""" - assert 0.0 <= c.min() and c.max() <= 1.0 - b,ch,h,w = c.shape - assert ch == 1 - - c = torch.nn.functional.interpolate(c, scale_factor=1/self.down_factor, - mode="area") - c = c.clamp(0.0, 1.0) - c = self.n_embed*c - c_quant = c.round() - c_ind = c_quant.to(dtype=torch.long) - - info = None, None, c_ind - return c_quant, None, info - - def decode(self, c): - c = c/self.n_embed - c = torch.nn.functional.interpolate(c, scale_factor=self.down_factor, - mode="nearest") - return c diff --git a/taming/modules/transformer/mingpt.py b/taming/modules/transformer/mingpt.py deleted file mode 100644 index d14b7b6..0000000 --- a/taming/modules/transformer/mingpt.py +++ /dev/null @@ -1,415 +0,0 @@ -""" -taken from: https://github.com/karpathy/minGPT/ -GPT model: -- the initial stem consists of a combination of token encoding and a positional encoding -- the meat of it is a uniform sequence of Transformer blocks - - each Transformer is a sequential combination of a 1-hidden-layer MLP block and a self-attention block - - all blocks feed into a central residual pathway similar to resnets -- the final decoder is a linear projection into a vanilla Softmax classifier -""" - -import math -import logging - -import torch -import torch.nn as nn -from torch.nn import functional as F -from transformers import top_k_top_p_filtering - -logger = logging.getLogger(__name__) - - -class GPTConfig: - """ base GPT config, params common to all GPT versions """ - embd_pdrop = 0.1 - resid_pdrop = 0.1 - attn_pdrop = 0.1 - - def __init__(self, vocab_size, block_size, **kwargs): - self.vocab_size = vocab_size - self.block_size = block_size - for k,v in kwargs.items(): - setattr(self, k, v) - - -class GPT1Config(GPTConfig): - """ GPT-1 like network roughly 125M params """ - n_layer = 12 - n_head = 12 - n_embd = 768 - - -class CausalSelfAttention(nn.Module): - """ - A vanilla multi-head masked self-attention layer with a projection at the end. - It is possible to use torch.nn.MultiheadAttention here but I am including an - explicit implementation here to show that there is nothing too scary here. - """ - - def __init__(self, config): - super().__init__() - assert config.n_embd % config.n_head == 0 - # key, query, value projections for all heads - self.key = nn.Linear(config.n_embd, config.n_embd) - self.query = nn.Linear(config.n_embd, config.n_embd) - self.value = nn.Linear(config.n_embd, config.n_embd) - # regularization - self.attn_drop = nn.Dropout(config.attn_pdrop) - self.resid_drop = nn.Dropout(config.resid_pdrop) - # output projection - self.proj = nn.Linear(config.n_embd, config.n_embd) - # causal mask to ensure that attention is only applied to the left in the input sequence - mask = torch.tril(torch.ones(config.block_size, - config.block_size)) - if hasattr(config, "n_unmasked"): - mask[:config.n_unmasked, :config.n_unmasked] = 1 - self.register_buffer("mask", mask.view(1, 1, config.block_size, config.block_size)) - self.n_head = config.n_head - - def forward(self, x, layer_past=None): - B, T, C = x.size() - - # calculate query, key, values for all heads in batch and move head forward to be the batch dim - k = self.key(x).view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs) - q = self.query(x).view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs) - v = self.value(x).view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs) - - present = torch.stack((k, v)) - if layer_past is not None: - past_key, past_value = layer_past - k = torch.cat((past_key, k), dim=-2) - v = torch.cat((past_value, v), dim=-2) - - # causal self-attention; Self-attend: (B, nh, T, hs) x (B, nh, hs, T) -> (B, nh, T, T) - att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1))) - if layer_past is None: - att = att.masked_fill(self.mask[:,:,:T,:T] == 0, float('-inf')) - - att = F.softmax(att, dim=-1) - att = self.attn_drop(att) - y = att @ v # (B, nh, T, T) x (B, nh, T, hs) -> (B, nh, T, hs) - y = y.transpose(1, 2).contiguous().view(B, T, C) # re-assemble all head outputs side by side - - # output projection - y = self.resid_drop(self.proj(y)) - return y, present # TODO: check that this does not break anything - - -class Block(nn.Module): - """ an unassuming Transformer block """ - def __init__(self, config): - super().__init__() - self.ln1 = nn.LayerNorm(config.n_embd) - self.ln2 = nn.LayerNorm(config.n_embd) - self.attn = CausalSelfAttention(config) - self.mlp = nn.Sequential( - nn.Linear(config.n_embd, 4 * config.n_embd), - nn.GELU(), # nice - nn.Linear(4 * config.n_embd, config.n_embd), - nn.Dropout(config.resid_pdrop), - ) - - def forward(self, x, layer_past=None, return_present=False): - # TODO: check that training still works - if return_present: assert not self.training - # layer past: tuple of length two with B, nh, T, hs - attn, present = self.attn(self.ln1(x), layer_past=layer_past) - - x = x + attn - x = x + self.mlp(self.ln2(x)) - if layer_past is not None or return_present: - return x, present - return x - - -class GPT(nn.Module): - """ the full GPT language model, with a context size of block_size """ - def __init__(self, vocab_size, block_size, n_layer=12, n_head=8, n_embd=256, - embd_pdrop=0., resid_pdrop=0., attn_pdrop=0., n_unmasked=0): - super().__init__() - config = GPTConfig(vocab_size=vocab_size, block_size=block_size, - embd_pdrop=embd_pdrop, resid_pdrop=resid_pdrop, attn_pdrop=attn_pdrop, - n_layer=n_layer, n_head=n_head, n_embd=n_embd, - n_unmasked=n_unmasked) - # input embedding stem - self.tok_emb = nn.Embedding(config.vocab_size, config.n_embd) - self.pos_emb = nn.Parameter(torch.zeros(1, config.block_size, config.n_embd)) - self.drop = nn.Dropout(config.embd_pdrop) - # transformer - self.blocks = nn.Sequential(*[Block(config) for _ in range(config.n_layer)]) - # decoder head - self.ln_f = nn.LayerNorm(config.n_embd) - self.head = nn.Linear(config.n_embd, config.vocab_size, bias=False) - self.block_size = config.block_size - self.apply(self._init_weights) - self.config = config - logger.info("number of parameters: %e", sum(p.numel() for p in self.parameters())) - - def get_block_size(self): - return self.block_size - - def _init_weights(self, module): - if isinstance(module, (nn.Linear, nn.Embedding)): - module.weight.data.normal_(mean=0.0, std=0.02) - if isinstance(module, nn.Linear) and module.bias is not None: - module.bias.data.zero_() - elif isinstance(module, nn.LayerNorm): - module.bias.data.zero_() - module.weight.data.fill_(1.0) - - def forward(self, idx, embeddings=None, targets=None): - # forward the GPT model - token_embeddings = self.tok_emb(idx) # each index maps to a (learnable) vector - - if embeddings is not None: # prepend explicit embeddings - token_embeddings = torch.cat((embeddings, token_embeddings), dim=1) - - t = token_embeddings.shape[1] - assert t <= self.block_size, "Cannot forward, model block size is exhausted." - position_embeddings = self.pos_emb[:, :t, :] # each position maps to a (learnable) vector - x = self.drop(token_embeddings + position_embeddings) - x = self.blocks(x) - x = self.ln_f(x) - logits = self.head(x) - - # if we are given some desired targets also calculate the loss - loss = None - if targets is not None: - loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1)) - - return logits, loss - - def forward_with_past(self, idx, embeddings=None, targets=None, past=None, past_length=None): - # inference only - assert not self.training - token_embeddings = self.tok_emb(idx) # each index maps to a (learnable) vector - if embeddings is not None: # prepend explicit embeddings - token_embeddings = torch.cat((embeddings, token_embeddings), dim=1) - - if past is not None: - assert past_length is not None - past = torch.cat(past, dim=-2) # n_layer, 2, b, nh, len_past, dim_head - past_shape = list(past.shape) - expected_shape = [self.config.n_layer, 2, idx.shape[0], self.config.n_head, past_length, self.config.n_embd//self.config.n_head] - assert past_shape == expected_shape, f"{past_shape} =/= {expected_shape}" - position_embeddings = self.pos_emb[:, past_length, :] # each position maps to a (learnable) vector - else: - position_embeddings = self.pos_emb[:, :token_embeddings.shape[1], :] - - x = self.drop(token_embeddings + position_embeddings) - presents = [] # accumulate over layers - for i, block in enumerate(self.blocks): - x, present = block(x, layer_past=past[i, ...] if past is not None else None, return_present=True) - presents.append(present) - - x = self.ln_f(x) - logits = self.head(x) - # if we are given some desired targets also calculate the loss - loss = None - if targets is not None: - loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1)) - - return logits, loss, torch.stack(presents) # _, _, n_layer, 2, b, nh, 1, dim_head - - -class DummyGPT(nn.Module): - # for debugging - def __init__(self, add_value=1): - super().__init__() - self.add_value = add_value - - def forward(self, idx): - return idx + self.add_value, None - - -class CodeGPT(nn.Module): - """Takes in semi-embeddings""" - def __init__(self, vocab_size, block_size, in_channels, n_layer=12, n_head=8, n_embd=256, - embd_pdrop=0., resid_pdrop=0., attn_pdrop=0., n_unmasked=0): - super().__init__() - config = GPTConfig(vocab_size=vocab_size, block_size=block_size, - embd_pdrop=embd_pdrop, resid_pdrop=resid_pdrop, attn_pdrop=attn_pdrop, - n_layer=n_layer, n_head=n_head, n_embd=n_embd, - n_unmasked=n_unmasked) - # input embedding stem - self.tok_emb = nn.Linear(in_channels, config.n_embd) - self.pos_emb = nn.Parameter(torch.zeros(1, config.block_size, config.n_embd)) - self.drop = nn.Dropout(config.embd_pdrop) - # transformer - self.blocks = nn.Sequential(*[Block(config) for _ in range(config.n_layer)]) - # decoder head - self.ln_f = nn.LayerNorm(config.n_embd) - self.head = nn.Linear(config.n_embd, config.vocab_size, bias=False) - self.block_size = config.block_size - self.apply(self._init_weights) - self.config = config - logger.info("number of parameters: %e", sum(p.numel() for p in self.parameters())) - - def get_block_size(self): - return self.block_size - - def _init_weights(self, module): - if isinstance(module, (nn.Linear, nn.Embedding)): - module.weight.data.normal_(mean=0.0, std=0.02) - if isinstance(module, nn.Linear) and module.bias is not None: - module.bias.data.zero_() - elif isinstance(module, nn.LayerNorm): - module.bias.data.zero_() - module.weight.data.fill_(1.0) - - def forward(self, idx, embeddings=None, targets=None): - # forward the GPT model - token_embeddings = self.tok_emb(idx) # each index maps to a (learnable) vector - - if embeddings is not None: # prepend explicit embeddings - token_embeddings = torch.cat((embeddings, token_embeddings), dim=1) - - t = token_embeddings.shape[1] - assert t <= self.block_size, "Cannot forward, model block size is exhausted." - position_embeddings = self.pos_emb[:, :t, :] # each position maps to a (learnable) vector - x = self.drop(token_embeddings + position_embeddings) - x = self.blocks(x) - x = self.taming_cinln_f(x) - logits = self.head(x) - - # if we are given some desired targets also calculate the loss - loss = None - if targets is not None: - loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1)) - - return logits, loss - - - -#### sampling utils - -def top_k_logits(logits, k): - v, ix = torch.topk(logits, k) - out = logits.clone() - out[out < v[:, [-1]]] = -float('Inf') - return out - -@torch.no_grad() -def sample(model, x, steps, temperature=1.0, sample=False, top_k=None): - """ - take a conditioning sequence of indices in x (of shape (b,t)) and predict the next token in - the sequence, feeding the predictions back into the model each time. Clearly the sampling - has quadratic complexity unlike an RNN that is only linear, and has a finite context window - of block_size, unlike an RNN that has an infinite context window. - """ - block_size = model.get_block_size() - model.eval() - for k in range(steps): - x_cond = x if x.size(1) <= block_size else x[:, -block_size:] # crop context if needed - logits, _ = model(x_cond) - # pluck the logits at the final step and scale by temperature - logits = logits[:, -1, :] / temperature - # optionally crop probabilities to only the top k options - if top_k is not None: - logits = top_k_logits(logits, top_k) - # apply softmax to convert to probabilities - probs = F.softmax(logits, dim=-1) - # sample from the distribution or take the most likely - if sample: - ix = torch.multinomial(probs, num_samples=1) - else: - _, ix = torch.topk(probs, k=1, dim=-1) - # append to the sequence and continue - x = torch.cat((x, ix), dim=1) - - return x - - -@torch.no_grad() -def sample_with_past(x, model, steps, temperature=1., sample_logits=True, - top_k=None, top_p=None, callback=None): - # x is conditioning - sample = x - cond_len = x.shape[1] - past = None - for n in range(steps): - if callback is not None: - callback(n) - logits, _, present = model.forward_with_past(x, past=past, past_length=(n+cond_len-1)) - if past is None: - past = [present] - else: - past.append(present) - logits = logits[:, -1, :] / temperature - if top_k is not None: - logits = top_k_top_p_filtering(logits, top_k=top_k, top_p=top_p) - - probs = F.softmax(logits, dim=-1) - if not sample_logits: - _, x = torch.topk(probs, k=1, dim=-1) - else: - x = torch.multinomial(probs, num_samples=1) - # append to the sequence and continue - sample = torch.cat((sample, x), dim=1) - del past - sample = sample[:, cond_len:] # cut conditioning off - return sample - - -#### clustering utils - -class KMeans(nn.Module): - def __init__(self, ncluster=512, nc=3, niter=10): - super().__init__() - self.ncluster = ncluster - self.nc = nc - self.niter = niter - self.shape = (3,32,32) - self.register_buffer("C", torch.zeros(self.ncluster,nc)) - self.register_buffer('initialized', torch.tensor(0, dtype=torch.uint8)) - - def is_initialized(self): - return self.initialized.item() == 1 - - @torch.no_grad() - def initialize(self, x): - N, D = x.shape - assert D == self.nc, D - c = x[torch.randperm(N)[:self.ncluster]] # init clusters at random - for i in range(self.niter): - # assign all pixels to the closest codebook element - a = ((x[:, None, :] - c[None, :, :])**2).sum(-1).argmin(1) - # move each codebook element to be the mean of the pixels that assigned to it - c = torch.stack([x[a==k].mean(0) for k in range(self.ncluster)]) - # re-assign any poorly positioned codebook elements - nanix = torch.any(torch.isnan(c), dim=1) - ndead = nanix.sum().item() - print('done step %d/%d, re-initialized %d dead clusters' % (i+1, self.niter, ndead)) - c[nanix] = x[torch.randperm(N)[:ndead]] # re-init dead clusters - - self.C.copy_(c) - self.initialized.fill_(1) - - - def forward(self, x, reverse=False, shape=None): - if not reverse: - # flatten - bs,c,h,w = x.shape - assert c == self.nc - x = x.reshape(bs,c,h*w,1) - C = self.C.permute(1,0) - C = C.reshape(1,c,1,self.ncluster) - a = ((x-C)**2).sum(1).argmin(-1) # bs, h*w indices - return a - else: - # flatten - bs, HW = x.shape - """ - c = self.C.reshape( 1, self.nc, 1, self.ncluster) - c = c[bs*[0],:,:,:] - c = c[:,:,HW*[0],:] - x = x.reshape(bs, 1, HW, 1) - x = x[:,3*[0],:,:] - x = torch.gather(c, dim=3, index=x) - """ - x = self.C[x] - x = x.permute(0,2,1) - shape = shape if shape is not None else self.shape - x = x.reshape(bs, *shape) - - return x diff --git a/taming/modules/transformer/permuter.py b/taming/modules/transformer/permuter.py deleted file mode 100644 index 0d43bb1..0000000 --- a/taming/modules/transformer/permuter.py +++ /dev/null @@ -1,248 +0,0 @@ -import torch -import torch.nn as nn -import numpy as np - - -class AbstractPermuter(nn.Module): - def __init__(self, *args, **kwargs): - super().__init__() - def forward(self, x, reverse=False): - raise NotImplementedError - - -class Identity(AbstractPermuter): - def __init__(self): - super().__init__() - - def forward(self, x, reverse=False): - return x - - -class Subsample(AbstractPermuter): - def __init__(self, H, W): - super().__init__() - C = 1 - indices = np.arange(H*W).reshape(C,H,W) - while min(H, W) > 1: - indices = indices.reshape(C,H//2,2,W//2,2) - indices = indices.transpose(0,2,4,1,3) - indices = indices.reshape(C*4,H//2, W//2) - H = H//2 - W = W//2 - C = C*4 - assert H == W == 1 - idx = torch.tensor(indices.ravel()) - self.register_buffer('forward_shuffle_idx', - nn.Parameter(idx, requires_grad=False)) - self.register_buffer('backward_shuffle_idx', - nn.Parameter(torch.argsort(idx), requires_grad=False)) - - def forward(self, x, reverse=False): - if not reverse: - return x[:, self.forward_shuffle_idx] - else: - return x[:, self.backward_shuffle_idx] - - -def mortonify(i, j): - """(i,j) index to linear morton code""" - i = np.uint64(i) - j = np.uint64(j) - - z = np.uint(0) - - for pos in range(32): - z = (z | - ((j & (np.uint64(1) << np.uint64(pos))) << np.uint64(pos)) | - ((i & (np.uint64(1) << np.uint64(pos))) << np.uint64(pos+1)) - ) - return z - - -class ZCurve(AbstractPermuter): - def __init__(self, H, W): - super().__init__() - reverseidx = [np.int64(mortonify(i,j)) for i in range(H) for j in range(W)] - idx = np.argsort(reverseidx) - idx = torch.tensor(idx) - reverseidx = torch.tensor(reverseidx) - self.register_buffer('forward_shuffle_idx', - idx) - self.register_buffer('backward_shuffle_idx', - reverseidx) - - def forward(self, x, reverse=False): - if not reverse: - return x[:, self.forward_shuffle_idx] - else: - return x[:, self.backward_shuffle_idx] - - -class SpiralOut(AbstractPermuter): - def __init__(self, H, W): - super().__init__() - assert H == W - size = W - indices = np.arange(size*size).reshape(size,size) - - i0 = size//2 - j0 = size//2-1 - - i = i0 - j = j0 - - idx = [indices[i0, j0]] - step_mult = 0 - for c in range(1, size//2+1): - step_mult += 1 - # steps left - for k in range(step_mult): - i = i - 1 - j = j - idx.append(indices[i, j]) - - # step down - for k in range(step_mult): - i = i - j = j + 1 - idx.append(indices[i, j]) - - step_mult += 1 - if c < size//2: - # step right - for k in range(step_mult): - i = i + 1 - j = j - idx.append(indices[i, j]) - - # step up - for k in range(step_mult): - i = i - j = j - 1 - idx.append(indices[i, j]) - else: - # end reached - for k in range(step_mult-1): - i = i + 1 - idx.append(indices[i, j]) - - assert len(idx) == size*size - idx = torch.tensor(idx) - self.register_buffer('forward_shuffle_idx', idx) - self.register_buffer('backward_shuffle_idx', torch.argsort(idx)) - - def forward(self, x, reverse=False): - if not reverse: - return x[:, self.forward_shuffle_idx] - else: - return x[:, self.backward_shuffle_idx] - - -class SpiralIn(AbstractPermuter): - def __init__(self, H, W): - super().__init__() - assert H == W - size = W - indices = np.arange(size*size).reshape(size,size) - - i0 = size//2 - j0 = size//2-1 - - i = i0 - j = j0 - - idx = [indices[i0, j0]] - step_mult = 0 - for c in range(1, size//2+1): - step_mult += 1 - # steps left - for k in range(step_mult): - i = i - 1 - j = j - idx.append(indices[i, j]) - - # step down - for k in range(step_mult): - i = i - j = j + 1 - idx.append(indices[i, j]) - - step_mult += 1 - if c < size//2: - # step right - for k in range(step_mult): - i = i + 1 - j = j - idx.append(indices[i, j]) - - # step up - for k in range(step_mult): - i = i - j = j - 1 - idx.append(indices[i, j]) - else: - # end reached - for k in range(step_mult-1): - i = i + 1 - idx.append(indices[i, j]) - - assert len(idx) == size*size - idx = idx[::-1] - idx = torch.tensor(idx) - self.register_buffer('forward_shuffle_idx', idx) - self.register_buffer('backward_shuffle_idx', torch.argsort(idx)) - - def forward(self, x, reverse=False): - if not reverse: - return x[:, self.forward_shuffle_idx] - else: - return x[:, self.backward_shuffle_idx] - - -class Random(nn.Module): - def __init__(self, H, W): - super().__init__() - indices = np.random.RandomState(1).permutation(H*W) - idx = torch.tensor(indices.ravel()) - self.register_buffer('forward_shuffle_idx', idx) - self.register_buffer('backward_shuffle_idx', torch.argsort(idx)) - - def forward(self, x, reverse=False): - if not reverse: - return x[:, self.forward_shuffle_idx] - else: - return x[:, self.backward_shuffle_idx] - - -class AlternateParsing(AbstractPermuter): - def __init__(self, H, W): - super().__init__() - indices = np.arange(W*H).reshape(H,W) - for i in range(1, H, 2): - indices[i, :] = indices[i, ::-1] - idx = indices.flatten() - assert len(idx) == H*W - idx = torch.tensor(idx) - self.register_buffer('forward_shuffle_idx', idx) - self.register_buffer('backward_shuffle_idx', torch.argsort(idx)) - - def forward(self, x, reverse=False): - if not reverse: - return x[:, self.forward_shuffle_idx] - else: - return x[:, self.backward_shuffle_idx] - - -if __name__ == "__main__": - p0 = AlternateParsing(16, 16) - print(p0.forward_shuffle_idx) - print(p0.backward_shuffle_idx) - - x = torch.randint(0, 768, size=(11, 256)) - y = p0(x) - xre = p0(y, reverse=True) - assert torch.equal(x, xre) - - p1 = SpiralOut(2, 2) - print(p1.forward_shuffle_idx) - print(p1.backward_shuffle_idx) diff --git a/taming/modules/util.py b/taming/modules/util.py deleted file mode 100644 index 9ee1638..0000000 --- a/taming/modules/util.py +++ /dev/null @@ -1,130 +0,0 @@ -import torch -import torch.nn as nn - - -def count_params(model): - total_params = sum(p.numel() for p in model.parameters()) - return total_params - - -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 - - -class AbstractEncoder(nn.Module): - def __init__(self): - super().__init__() - - def encode(self, *args, **kwargs): - raise NotImplementedError - - -class Labelator(AbstractEncoder): - """Net2Net Interface for Class-Conditional Model""" - def __init__(self, n_classes, quantize_interface=True): - super().__init__() - self.n_classes = n_classes - self.quantize_interface = quantize_interface - - def encode(self, c): - c = c[:,None] - if self.quantize_interface: - return c, None, [None, None, c.long()] - return c - - -class SOSProvider(AbstractEncoder): - # for unconditional training - def __init__(self, sos_token, quantize_interface=True): - super().__init__() - self.sos_token = sos_token - self.quantize_interface = quantize_interface - - def encode(self, x): - # get batch size from data and replicate sos_token - c = torch.ones(x.shape[0], 1)*self.sos_token - c = c.long().to(x.device) - if self.quantize_interface: - return c, None, [None, None, c] - return c diff --git a/taming/modules/vqvae/quantize.py b/taming/modules/vqvae/quantize.py deleted file mode 100644 index d75544e..0000000 --- a/taming/modules/vqvae/quantize.py +++ /dev/null @@ -1,445 +0,0 @@ -import torch -import torch.nn as nn -import torch.nn.functional as F -import numpy as np -from torch import einsum -from einops import rearrange - - -class VectorQuantizer(nn.Module): - """ - see https://github.com/MishaLaskin/vqvae/blob/d761a999e2267766400dc646d82d3ac3657771d4/models/quantizer.py - ____________________________________________ - Discretization bottleneck part of the VQ-VAE. - Inputs: - - n_e : number of embeddings - - e_dim : dimension of embedding - - beta : commitment cost used in loss term, beta * ||z_e(x)-sg[e]||^2 - _____________________________________________ - """ - - # NOTE: this class contains a bug regarding beta; see VectorQuantizer2 for - # a fix and use legacy=False to apply that fix. VectorQuantizer2 can be - # used wherever VectorQuantizer has been used before and is additionally - # more efficient. - def __init__(self, n_e, e_dim, beta): - super(VectorQuantizer, self).__init__() - self.n_e = n_e - self.e_dim = e_dim - self.beta = beta - - self.embedding = nn.Embedding(self.n_e, self.e_dim) - self.embedding.weight.data.uniform_(-1.0 / self.n_e, 1.0 / self.n_e) - - def forward(self, z): - """ - Inputs the output of the encoder network z and maps it to a discrete - one-hot vector that is the index of the closest embedding vector e_j - z (continuous) -> z_q (discrete) - z.shape = (batch, channel, height, width) - quantization pipeline: - 1. get encoder input (B,C,H,W) - 2. flatten input to (B*H*W,C) - """ - # reshape z -> (batch, height, width, channel) and flatten - z = z.permute(0, 2, 3, 1).contiguous() - z_flattened = z.view(-1, self.e_dim) - # distances from z to embeddings e_j (z - e)^2 = z^2 + e^2 - 2 e * z - - d = torch.sum(z_flattened ** 2, dim=1, keepdim=True) + \ - torch.sum(self.embedding.weight**2, dim=1) - 2 * \ - torch.matmul(z_flattened, self.embedding.weight.t()) - - ## could possible replace this here - # #\start... - # find closest encodings - min_encoding_indices = torch.argmin(d, dim=1).unsqueeze(1) - - min_encodings = torch.zeros( - min_encoding_indices.shape[0], self.n_e).to(z) - min_encodings.scatter_(1, min_encoding_indices, 1) - - # dtype min encodings: torch.float32 - # min_encodings shape: torch.Size([2048, 512]) - # min_encoding_indices.shape: torch.Size([2048, 1]) - - # get quantized latent vectors - z_q = torch.matmul(min_encodings, self.embedding.weight).view(z.shape) - #.........\end - - # with: - # .........\start - #min_encoding_indices = torch.argmin(d, dim=1) - #z_q = self.embedding(min_encoding_indices) - # ......\end......... (TODO) - - # compute loss for embedding - loss = torch.mean((z_q.detach()-z)**2) + self.beta * \ - torch.mean((z_q - z.detach()) ** 2) - - # preserve gradients - z_q = z + (z_q - z).detach() - - # perplexity - e_mean = torch.mean(min_encodings, dim=0) - perplexity = torch.exp(-torch.sum(e_mean * torch.log(e_mean + 1e-10))) - - # reshape back to match original input shape - z_q = z_q.permute(0, 3, 1, 2).contiguous() - - return z_q, loss, (perplexity, min_encodings, min_encoding_indices) - - def get_codebook_entry(self, indices, shape): - # shape specifying (batch, height, width, channel) - # TODO: check for more easy handling with nn.Embedding - min_encodings = torch.zeros(indices.shape[0], self.n_e).to(indices) - min_encodings.scatter_(1, indices[:,None], 1) - - # get quantized latent vectors - z_q = torch.matmul(min_encodings.float(), self.embedding.weight) - - if shape is not None: - z_q = z_q.view(shape) - - # reshape back to match original input shape - z_q = z_q.permute(0, 3, 1, 2).contiguous() - - return z_q - - -class GumbelQuantize(nn.Module): - """ - credit to @karpathy: https://github.com/karpathy/deep-vector-quantization/blob/main/model.py (thanks!) - Gumbel Softmax trick quantizer - Categorical Reparameterization with Gumbel-Softmax, Jang et al. 2016 - https://arxiv.org/abs/1611.01144 - """ - def __init__(self, num_hiddens, embedding_dim, n_embed, straight_through=True, - kl_weight=5e-4, temp_init=1.0, use_vqinterface=True, - remap=None, unknown_index="random"): - super().__init__() - - self.embedding_dim = embedding_dim - self.n_embed = n_embed - - self.straight_through = straight_through - self.temperature = temp_init - self.kl_weight = kl_weight - - self.proj = nn.Conv2d(num_hiddens, n_embed, 1) - self.embed = nn.Embedding(n_embed, embedding_dim) - - self.use_vqinterface = use_vqinterface - - self.remap = remap - if self.remap is not None: - self.register_buffer("used", torch.tensor(np.load(self.remap))) - self.re_embed = self.used.shape[0] - self.unknown_index = unknown_index # "random" or "extra" or integer - if self.unknown_index == "extra": - self.unknown_index = self.re_embed - self.re_embed = self.re_embed+1 - print(f"Remapping {self.n_embed} indices to {self.re_embed} indices. " - f"Using {self.unknown_index} for unknown indices.") - else: - self.re_embed = n_embed - - def remap_to_used(self, inds): - ishape = inds.shape - assert len(ishape)>1 - inds = inds.reshape(ishape[0],-1) - used = self.used.to(inds) - match = (inds[:,:,None]==used[None,None,...]).long() - new = match.argmax(-1) - unknown = match.sum(2)<1 - if self.unknown_index == "random": - new[unknown]=torch.randint(0,self.re_embed,size=new[unknown].shape).to(device=new.device) - else: - new[unknown] = self.unknown_index - return new.reshape(ishape) - - def unmap_to_all(self, inds): - ishape = inds.shape - assert len(ishape)>1 - inds = inds.reshape(ishape[0],-1) - used = self.used.to(inds) - if self.re_embed > self.used.shape[0]: # extra token - inds[inds>=self.used.shape[0]] = 0 # simply set to zero - back=torch.gather(used[None,:][inds.shape[0]*[0],:], 1, inds) - return back.reshape(ishape) - - def forward(self, z, temp=None, return_logits=False): - # force hard = True when we are in eval mode, as we must quantize. actually, always true seems to work - hard = self.straight_through if self.training else True - temp = self.temperature if temp is None else temp - - logits = self.proj(z) - if self.remap is not None: - # continue only with used logits - full_zeros = torch.zeros_like(logits) - logits = logits[:,self.used,...] - - soft_one_hot = F.gumbel_softmax(logits, tau=temp, dim=1, hard=hard) - if self.remap is not None: - # go back to all entries but unused set to zero - full_zeros[:,self.used,...] = soft_one_hot - soft_one_hot = full_zeros - z_q = einsum('b n h w, n d -> b d h w', soft_one_hot, self.embed.weight) - - # + kl divergence to the prior loss - qy = F.softmax(logits, dim=1) - diff = self.kl_weight * torch.sum(qy * torch.log(qy * self.n_embed + 1e-10), dim=1).mean() - - ind = soft_one_hot.argmax(dim=1) - if self.remap is not None: - ind = self.remap_to_used(ind) - if self.use_vqinterface: - if return_logits: - return z_q, diff, (None, None, ind), logits - return z_q, diff, (None, None, ind) - return z_q, diff, ind - - def get_codebook_entry(self, indices, shape): - b, h, w, c = shape - assert b*h*w == indices.shape[0] - indices = rearrange(indices, '(b h w) -> b h w', b=b, h=h, w=w) - if self.remap is not None: - indices = self.unmap_to_all(indices) - one_hot = F.one_hot(indices, num_classes=self.n_embed).permute(0, 3, 1, 2).float() - z_q = einsum('b n h w, n d -> b d h w', one_hot, self.embed.weight) - return z_q - - -class VectorQuantizer2(nn.Module): - """ - Improved version over VectorQuantizer, can be used as a drop-in replacement. Mostly - avoids costly matrix multiplications and allows for post-hoc remapping of indices. - """ - # NOTE: due to a bug the beta term was applied to the wrong term. for - # backwards compatibility we use the buggy version by default, but you can - # specify legacy=False to fix it. - def __init__(self, n_e, e_dim, beta, remap=None, unknown_index="random", - sane_index_shape=False, legacy=True): - super().__init__() - self.n_e = n_e - self.e_dim = e_dim - self.beta = beta - self.legacy = legacy - - self.embedding = nn.Embedding(self.n_e, self.e_dim) - self.embedding.weight.data.uniform_(-1.0 / self.n_e, 1.0 / self.n_e) - - self.remap = remap - if self.remap is not None: - self.register_buffer("used", torch.tensor(np.load(self.remap))) - self.re_embed = self.used.shape[0] - self.unknown_index = unknown_index # "random" or "extra" or integer - if self.unknown_index == "extra": - self.unknown_index = self.re_embed - self.re_embed = self.re_embed+1 - print(f"Remapping {self.n_e} indices to {self.re_embed} indices. " - f"Using {self.unknown_index} for unknown indices.") - else: - self.re_embed = n_e - - self.sane_index_shape = sane_index_shape - - def remap_to_used(self, inds): - ishape = inds.shape - assert len(ishape)>1 - inds = inds.reshape(ishape[0],-1) - used = self.used.to(inds) - match = (inds[:,:,None]==used[None,None,...]).long() - new = match.argmax(-1) - unknown = match.sum(2)<1 - if self.unknown_index == "random": - new[unknown]=torch.randint(0,self.re_embed,size=new[unknown].shape).to(device=new.device) - else: - new[unknown] = self.unknown_index - return new.reshape(ishape) - - def unmap_to_all(self, inds): - ishape = inds.shape - assert len(ishape)>1 - inds = inds.reshape(ishape[0],-1) - used = self.used.to(inds) - if self.re_embed > self.used.shape[0]: # extra token - inds[inds>=self.used.shape[0]] = 0 # simply set to zero - back=torch.gather(used[None,:][inds.shape[0]*[0],:], 1, inds) - return back.reshape(ishape) - - def forward(self, z, temp=None, rescale_logits=False, return_logits=False): - assert temp is None or temp==1.0, "Only for interface compatible with Gumbel" - assert rescale_logits==False, "Only for interface compatible with Gumbel" - assert return_logits==False, "Only for interface compatible with Gumbel" - # reshape z -> (batch, height, width, channel) and flatten - z = rearrange(z, 'b c h w -> b h w c').contiguous() - z_flattened = z.view(-1, self.e_dim) - # distances from z to embeddings e_j (z - e)^2 = z^2 + e^2 - 2 e * z - - d = torch.sum(z_flattened ** 2, dim=1, keepdim=True) + \ - torch.sum(self.embedding.weight**2, dim=1) - 2 * \ - torch.einsum('bd,dn->bn', z_flattened, rearrange(self.embedding.weight, 'n d -> d n')) - - min_encoding_indices = torch.argmin(d, dim=1) - z_q = self.embedding(min_encoding_indices).view(z.shape) - perplexity = None - min_encodings = None - - # compute loss for embedding - if not self.legacy: - loss = self.beta * torch.mean((z_q.detach()-z)**2) + \ - torch.mean((z_q - z.detach()) ** 2) - else: - loss = torch.mean((z_q.detach()-z)**2) + self.beta * \ - torch.mean((z_q - z.detach()) ** 2) - - # preserve gradients - z_q = z + (z_q - z).detach() - - # reshape back to match original input shape - z_q = rearrange(z_q, 'b h w c -> b c h w').contiguous() - - if self.remap is not None: - min_encoding_indices = min_encoding_indices.reshape(z.shape[0],-1) # add batch axis - min_encoding_indices = self.remap_to_used(min_encoding_indices) - min_encoding_indices = min_encoding_indices.reshape(-1,1) # flatten - - if self.sane_index_shape: - min_encoding_indices = min_encoding_indices.reshape( - z_q.shape[0], z_q.shape[2], z_q.shape[3]) - - return z_q, loss, (perplexity, min_encodings, min_encoding_indices) - - def get_codebook_entry(self, indices, shape): - # shape specifying (batch, height, width, channel) - if self.remap is not None: - indices = indices.reshape(shape[0],-1) # add batch axis - indices = self.unmap_to_all(indices) - indices = indices.reshape(-1) # flatten again - - # get quantized latent vectors - z_q = self.embedding(indices) - - if shape is not None: - z_q = z_q.view(shape) - # reshape back to match original input shape - z_q = z_q.permute(0, 3, 1, 2).contiguous() - - return z_q - -class EmbeddingEMA(nn.Module): - def __init__(self, num_tokens, codebook_dim, decay=0.99, eps=1e-5): - super().__init__() - self.decay = decay - self.eps = eps - weight = torch.randn(num_tokens, codebook_dim) - self.weight = nn.Parameter(weight, requires_grad = False) - self.cluster_size = nn.Parameter(torch.zeros(num_tokens), requires_grad = False) - self.embed_avg = nn.Parameter(weight.clone(), requires_grad = False) - self.update = True - - def forward(self, embed_id): - return F.embedding(embed_id, self.weight) - - def cluster_size_ema_update(self, new_cluster_size): - self.cluster_size.data.mul_(self.decay).add_(new_cluster_size, alpha=1 - self.decay) - - def embed_avg_ema_update(self, new_embed_avg): - self.embed_avg.data.mul_(self.decay).add_(new_embed_avg, alpha=1 - self.decay) - - def weight_update(self, num_tokens): - n = self.cluster_size.sum() - smoothed_cluster_size = ( - (self.cluster_size + self.eps) / (n + num_tokens * self.eps) * n - ) - #normalize embedding average with smoothed cluster size - embed_normalized = self.embed_avg / smoothed_cluster_size.unsqueeze(1) - self.weight.data.copy_(embed_normalized) - - -class EMAVectorQuantizer(nn.Module): - def __init__(self, n_embed, embedding_dim, beta, decay=0.99, eps=1e-5, - remap=None, unknown_index="random"): - super().__init__() - self.codebook_dim = codebook_dim - self.num_tokens = num_tokens - self.beta = beta - self.embedding = EmbeddingEMA(self.num_tokens, self.codebook_dim, decay, eps) - - self.remap = remap - if self.remap is not None: - self.register_buffer("used", torch.tensor(np.load(self.remap))) - self.re_embed = self.used.shape[0] - self.unknown_index = unknown_index # "random" or "extra" or integer - if self.unknown_index == "extra": - self.unknown_index = self.re_embed - self.re_embed = self.re_embed+1 - print(f"Remapping {self.n_embed} indices to {self.re_embed} indices. " - f"Using {self.unknown_index} for unknown indices.") - else: - self.re_embed = n_embed - - def remap_to_used(self, inds): - ishape = inds.shape - assert len(ishape)>1 - inds = inds.reshape(ishape[0],-1) - used = self.used.to(inds) - match = (inds[:,:,None]==used[None,None,...]).long() - new = match.argmax(-1) - unknown = match.sum(2)<1 - if self.unknown_index == "random": - new[unknown]=torch.randint(0,self.re_embed,size=new[unknown].shape).to(device=new.device) - else: - new[unknown] = self.unknown_index - return new.reshape(ishape) - - def unmap_to_all(self, inds): - ishape = inds.shape - assert len(ishape)>1 - inds = inds.reshape(ishape[0],-1) - used = self.used.to(inds) - if self.re_embed > self.used.shape[0]: # extra token - inds[inds>=self.used.shape[0]] = 0 # simply set to zero - back=torch.gather(used[None,:][inds.shape[0]*[0],:], 1, inds) - return back.reshape(ishape) - - def forward(self, z): - # reshape z -> (batch, height, width, channel) and flatten - #z, 'b c h w -> b h w c' - z = rearrange(z, 'b c h w -> b h w c') - z_flattened = z.reshape(-1, self.codebook_dim) - - # distances from z to embeddings e_j (z - e)^2 = z^2 + e^2 - 2 e * z - d = z_flattened.pow(2).sum(dim=1, keepdim=True) + \ - self.embedding.weight.pow(2).sum(dim=1) - 2 * \ - torch.einsum('bd,nd->bn', z_flattened, self.embedding.weight) # 'n d -> d n' - - - encoding_indices = torch.argmin(d, dim=1) - - z_q = self.embedding(encoding_indices).view(z.shape) - encodings = F.one_hot(encoding_indices, self.num_tokens).type(z.dtype) - avg_probs = torch.mean(encodings, dim=0) - perplexity = torch.exp(-torch.sum(avg_probs * torch.log(avg_probs + 1e-10))) - - if self.training and self.embedding.update: - #EMA cluster size - encodings_sum = encodings.sum(0) - self.embedding.cluster_size_ema_update(encodings_sum) - #EMA embedding average - embed_sum = encodings.transpose(0,1) @ z_flattened - self.embedding.embed_avg_ema_update(embed_sum) - #normalize embed_avg and update weight - self.embedding.weight_update(self.num_tokens) - - # compute loss for embedding - loss = self.beta * F.mse_loss(z_q.detach(), z) - - # preserve gradients - z_q = z + (z_q - z).detach() - - # reshape back to match original input shape - #z_q, 'b h w c -> b c h w' - z_q = rearrange(z_q, 'b h w c -> b c h w') - return z_q, loss, (perplexity, encodings, encoding_indices) diff --git a/taming/util.py b/taming/util.py deleted file mode 100644 index 06053e5..0000000 --- a/taming/util.py +++ /dev/null @@ -1,157 +0,0 @@ -import os, hashlib -import requests -from tqdm import tqdm - -URL_MAP = { - "vgg_lpips": "https://heibox.uni-heidelberg.de/f/607503859c864bc1b30b/?dl=1" -} - -CKPT_MAP = { - "vgg_lpips": "vgg.pth" -} - -MD5_MAP = { - "vgg_lpips": "d507d7349b931f0638a25a48a722f98a" -} - - -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 KeyNotFoundError(Exception): - def __init__(self, cause, keys=None, visited=None): - self.cause = cause - self.keys = keys - self.visited = visited - messages = list() - if keys is not None: - messages.append("Key not found: {}".format(keys)) - if visited is not None: - messages.append("Visited: {}".format(visited)) - messages.append("Cause:\n{}".format(cause)) - message = "\n".join(messages) - super().__init__(message) - - -def retrieve( - list_or_dict, key, splitval="/", default=None, expand=True, pass_success=False -): - """Given a nested list or dict return the desired value at key expanding - callable nodes if necessary and :attr:`expand` is ``True``. The expansion - is done in-place. - - Parameters - ---------- - list_or_dict : list or dict - Possibly nested list or dictionary. - key : str - key/to/value, path like string describing all keys necessary to - consider to get to the desired value. List indices can also be - passed here. - splitval : str - String that defines the delimiter between keys of the - different depth levels in `key`. - default : obj - Value returned if :attr:`key` is not found. - expand : bool - Whether to expand callable nodes on the path or not. - - Returns - ------- - The desired value or if :attr:`default` is not ``None`` and the - :attr:`key` is not found returns ``default``. - - Raises - ------ - Exception if ``key`` not in ``list_or_dict`` and :attr:`default` is - ``None``. - """ - - keys = key.split(splitval) - - success = True - try: - visited = [] - parent = None - last_key = None - for key in keys: - if callable(list_or_dict): - if not expand: - raise KeyNotFoundError( - ValueError( - "Trying to get past callable node with expand=False." - ), - keys=keys, - visited=visited, - ) - list_or_dict = list_or_dict() - parent[last_key] = list_or_dict - - last_key = key - parent = list_or_dict - - try: - if isinstance(list_or_dict, dict): - list_or_dict = list_or_dict[key] - else: - list_or_dict = list_or_dict[int(key)] - except (KeyError, IndexError, ValueError) as e: - raise KeyNotFoundError(e, keys=keys, visited=visited) - - visited += [key] - # final expansion of retrieved value - if expand and callable(list_or_dict): - list_or_dict = list_or_dict() - parent[last_key] = list_or_dict - except KeyNotFoundError as e: - if default is None: - raise e - else: - list_or_dict = default - success = False - - if not pass_success: - return list_or_dict - else: - return list_or_dict, success - - -if __name__ == "__main__": - config = {"keya": "a", - "keyb": "b", - "keyc": - {"cc1": 1, - "cc2": 2, - } - } - from omegaconf import OmegaConf - config = OmegaConf.create(config) - print(config) - retrieve(config, "keya") - diff --git a/utils/image/align_color.py b/utils/align_color.py similarity index 100% rename from utils/image/align_color.py rename to utils/align_color.py diff --git a/utils/image/__init__.py b/utils/image/__init__.py deleted file mode 100644 index 3bbe064..0000000 --- a/utils/image/__init__.py +++ /dev/null @@ -1,26 +0,0 @@ -from .diffjpeg import DiffJPEG -from .usm_sharp import USMSharp -from .common import ( - random_crop_arr, center_crop_arr, augment, - filter2D, rgb2ycbcr_pt, auto_resize, pad -) -from .align_color import ( - wavelet_reconstruction, adaptive_instance_normalization -) - -__all__ = [ - "DiffJPEG", - - "USMSharp", - - "random_crop_arr", - "center_crop_arr", - "augment", - "filter2D", - "rgb2ycbcr_pt", - "auto_resize", - "pad", - - "wavelet_reconstruction", - "adaptive_instance_normalization" -] diff --git a/utils/image/common.py b/utils/image/common.py deleted file mode 100644 index 2e8ca3c..0000000 --- a/utils/image/common.py +++ /dev/null @@ -1,235 +0,0 @@ -import random -import math - -from PIL import Image -import numpy as np -import cv2 -import torch -from torch.nn import functional as F - - -# https://github.com/openai/guided-diffusion/blob/main/guided_diffusion/image_datasets.py -def center_crop_arr(pil_image, image_size): - # We are not on a new enough PIL to support the `reducing_gap` - # argument, which uses BOX downsampling at powers of two first. - # Thus, we do it by hand to improve downsample quality. - while min(*pil_image.size) >= 2 * image_size: - pil_image = pil_image.resize( - tuple(x // 2 for x in pil_image.size), resample=Image.BOX - ) - - scale = image_size / min(*pil_image.size) - pil_image = pil_image.resize( - tuple(round(x * scale) for x in pil_image.size), resample=Image.BICUBIC - ) - - arr = np.array(pil_image) - crop_y = (arr.shape[0] - image_size) // 2 - crop_x = (arr.shape[1] - image_size) // 2 - return arr[crop_y : crop_y + image_size, crop_x : crop_x + image_size] - - -# https://github.com/openai/guided-diffusion/blob/main/guided_diffusion/image_datasets.py -def random_crop_arr(pil_image, image_size, min_crop_frac=0.8, max_crop_frac=1.0): - min_smaller_dim_size = math.ceil(image_size / max_crop_frac) - max_smaller_dim_size = math.ceil(image_size / min_crop_frac) - smaller_dim_size = random.randrange(min_smaller_dim_size, max_smaller_dim_size + 1) - - # We are not on a new enough PIL to support the `reducing_gap` - # argument, which uses BOX downsampling at powers of two first. - # Thus, we do it by hand to improve downsample quality. - while min(*pil_image.size) >= 2 * smaller_dim_size: - pil_image = pil_image.resize( - tuple(x // 2 for x in pil_image.size), resample=Image.BOX - ) - - scale = smaller_dim_size / min(*pil_image.size) - pil_image = pil_image.resize( - tuple(round(x * scale) for x in pil_image.size), resample=Image.BICUBIC - ) - - arr = np.array(pil_image) - crop_y = random.randrange(arr.shape[0] - image_size + 1) - crop_x = random.randrange(arr.shape[1] - image_size + 1) - return arr[crop_y : crop_y + image_size, crop_x : crop_x + image_size] - - -# https://github.com/XPixelGroup/BasicSR/blob/master/basicsr/data/transforms.py -def augment(imgs, hflip=True, rotation=True, flows=None, return_status=False): - """Augment: horizontal flips OR rotate (0, 90, 180, 270 degrees). - - We use vertical flip and transpose for rotation implementation. - All the images in the list use the same augmentation. - - Args: - imgs (list[ndarray] | ndarray): Images to be augmented. If the input - is an ndarray, it will be transformed to a list. - hflip (bool): Horizontal flip. Default: True. - rotation (bool): Ratotation. Default: True. - flows (list[ndarray]: Flows to be augmented. If the input is an - ndarray, it will be transformed to a list. - Dimension is (h, w, 2). Default: None. - return_status (bool): Return the status of flip and rotation. - Default: False. - - Returns: - list[ndarray] | ndarray: Augmented images and flows. If returned - results only have one element, just return ndarray. - - """ - hflip = hflip and random.random() < 0.5 - vflip = rotation and random.random() < 0.5 - rot90 = rotation and random.random() < 0.5 - - def _augment(img): - if hflip: # horizontal - cv2.flip(img, 1, img) - if vflip: # vertical - cv2.flip(img, 0, img) - if rot90: - img = img.transpose(1, 0, 2) - return img - - def _augment_flow(flow): - if hflip: # horizontal - cv2.flip(flow, 1, flow) - flow[:, :, 0] *= -1 - if vflip: # vertical - cv2.flip(flow, 0, flow) - flow[:, :, 1] *= -1 - if rot90: - flow = flow.transpose(1, 0, 2) - flow = flow[:, :, [1, 0]] - return flow - - if not isinstance(imgs, list): - imgs = [imgs] - imgs = [_augment(img) for img in imgs] - if len(imgs) == 1: - imgs = imgs[0] - - if flows is not None: - if not isinstance(flows, list): - flows = [flows] - flows = [_augment_flow(flow) for flow in flows] - if len(flows) == 1: - flows = flows[0] - return imgs, flows - else: - if return_status: - return imgs, (hflip, vflip, rot90) - else: - return imgs - - -# https://github.com/XPixelGroup/BasicSR/blob/master/basicsr/utils/img_process_util.py -def filter2D(img, kernel): - """PyTorch version of cv2.filter2D - - Args: - img (Tensor): (b, c, h, w) - kernel (Tensor): (b, k, k) - """ - k = kernel.size(-1) - b, c, h, w = img.size() - if k % 2 == 1: - img = F.pad(img, (k // 2, k // 2, k // 2, k // 2), mode='reflect') - else: - raise ValueError('Wrong kernel size') - - ph, pw = img.size()[-2:] - - if kernel.size(0) == 1: - # apply the same kernel to all batch images - img = img.view(b * c, 1, ph, pw) - kernel = kernel.view(1, 1, k, k) - return F.conv2d(img, kernel, padding=0).view(b, c, h, w) - else: - img = img.view(1, b * c, ph, pw) - kernel = kernel.view(b, 1, k, k).repeat(1, c, 1, 1).view(b * c, 1, k, k) - return F.conv2d(img, kernel, groups=b * c).view(b, c, h, w) - - -# https://github.com/XPixelGroup/BasicSR/blob/033cd6896d898fdd3dcda32e3102a792efa1b8f4/basicsr/utils/color_util.py#L186 -def rgb2ycbcr_pt(img, y_only=False): - """Convert RGB images to YCbCr images (PyTorch version). - - It implements the ITU-R BT.601 conversion for standard-definition television. See more details in - https://en.wikipedia.org/wiki/YCbCr#ITU-R_BT.601_conversion. - - Args: - img (Tensor): Images with shape (n, 3, h, w), the range [0, 1], float, RGB format. - y_only (bool): Whether to only return Y channel. Default: False. - - Returns: - (Tensor): converted images with the shape (n, 3/1, h, w), the range [0, 1], float. - """ - if y_only: - weight = torch.tensor([[65.481], [128.553], [24.966]]).to(img) - out_img = torch.matmul(img.permute(0, 2, 3, 1), weight).permute(0, 3, 1, 2) + 16.0 - else: - weight = torch.tensor([[65.481, -37.797, 112.0], [128.553, -74.203, -93.786], [24.966, 112.0, -18.214]]).to(img) - bias = torch.tensor([16, 128, 128]).view(1, 3, 1, 1).to(img) - out_img = torch.matmul(img.permute(0, 2, 3, 1), weight).permute(0, 3, 1, 2) + bias - - out_img = out_img / 255. - return out_img - - -def to_pil_image(inputs, mem_order, val_range, channel_order): - # convert inputs to numpy array - if isinstance(inputs, torch.Tensor): - inputs = inputs.cpu().numpy() - assert isinstance(inputs, np.ndarray) - - # make sure that inputs is a 4-dimension array - if mem_order in ["hwc", "chw"]: - inputs = inputs[None, ...] - mem_order = f"n{mem_order}" - # to NHWC - if mem_order == "nchw": - inputs = inputs.transpose(0, 2, 3, 1) - # to RGB - if channel_order == "bgr": - inputs = inputs[..., ::-1].copy() - else: - assert channel_order == "rgb" - - if val_range == "0,1": - inputs = inputs * 255 - elif val_range == "-1,1": - inputs = (inputs + 1) * 127.5 - else: - assert val_range == "0,255" - - inputs = inputs.clip(0, 255).astype(np.uint8) - return [inputs[i] for i in range(len(inputs))] - - -def put_text(pil_img_arr, text): - cv_img = pil_img_arr[..., ::-1].copy() - cv2.putText(cv_img, text, (10, 35), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2) - return cv_img[..., ::-1].copy() - - -def auto_resize(img: Image.Image, size: int) -> Image.Image: - short_edge = min(img.size) - if short_edge < size: - r = size / short_edge - img = img.resize( - tuple(math.ceil(x * r) for x in img.size), Image.BICUBIC - ) - else: - # make a deep copy of this image for safety - img = img.copy() - return img - - -def pad(img: np.ndarray, scale: int) -> np.ndarray: - h, w = img.shape[:2] - ph = 0 if h % scale == 0 else math.ceil(h / scale) * scale - h - pw = 0 if w % scale == 0 else math.ceil(w / scale) * scale - w - return np.pad( - img, pad_width=((0, ph), (0, pw), (0, 0)), mode="constant", - constant_values=0 - ) diff --git a/utils/image/diffjpeg.py b/utils/image/diffjpeg.py deleted file mode 100644 index 3392821..0000000 --- a/utils/image/diffjpeg.py +++ /dev/null @@ -1,492 +0,0 @@ -# https://github.com/XPixelGroup/BasicSR/blob/master/basicsr/utils/diffjpeg.py -""" -Modified from https://github.com/mlomnitz/DiffJPEG - -For images not divisible by 8 -https://dsp.stackexchange.com/questions/35339/jpeg-dct-padding/35343#35343 -""" -import itertools -import numpy as np -import torch -import torch.nn as nn -from torch.nn import functional as F - -# ------------------------ utils ------------------------# -y_table = np.array( - [[16, 11, 10, 16, 24, 40, 51, 61], [12, 12, 14, 19, 26, 58, 60, 55], [14, 13, 16, 24, 40, 57, 69, 56], - [14, 17, 22, 29, 51, 87, 80, 62], [18, 22, 37, 56, 68, 109, 103, 77], [24, 35, 55, 64, 81, 104, 113, 92], - [49, 64, 78, 87, 103, 121, 120, 101], [72, 92, 95, 98, 112, 100, 103, 99]], - dtype=np.float32).T -y_table = nn.Parameter(torch.from_numpy(y_table)) -c_table = np.empty((8, 8), dtype=np.float32) -c_table.fill(99) -c_table[:4, :4] = np.array([[17, 18, 24, 47], [18, 21, 26, 66], [24, 26, 56, 99], [47, 66, 99, 99]]).T -c_table = nn.Parameter(torch.from_numpy(c_table)) - - -def diff_round(x): - """ Differentiable rounding function - """ - return torch.round(x) + (x - torch.round(x))**3 - - -def quality_to_factor(quality): - """ Calculate factor corresponding to quality - - Args: - quality(float): Quality for jpeg compression. - - Returns: - float: Compression factor. - """ - if quality < 50: - quality = 5000. / quality - else: - quality = 200. - quality * 2 - return quality / 100. - - -# ------------------------ compression ------------------------# -class RGB2YCbCrJpeg(nn.Module): - """ Converts RGB image to YCbCr - """ - - def __init__(self): - super(RGB2YCbCrJpeg, self).__init__() - matrix = np.array([[0.299, 0.587, 0.114], [-0.168736, -0.331264, 0.5], [0.5, -0.418688, -0.081312]], - dtype=np.float32).T - self.shift = nn.Parameter(torch.tensor([0., 128., 128.])) - self.matrix = nn.Parameter(torch.from_numpy(matrix)) - - def forward(self, image): - """ - Args: - image(Tensor): batch x 3 x height x width - - Returns: - Tensor: batch x height x width x 3 - """ - image = image.permute(0, 2, 3, 1) - result = torch.tensordot(image, self.matrix, dims=1) + self.shift - return result.view(image.shape) - - -class ChromaSubsampling(nn.Module): - """ Chroma subsampling on CbCr channels - """ - - def __init__(self): - super(ChromaSubsampling, self).__init__() - - def forward(self, image): - """ - Args: - image(tensor): batch x height x width x 3 - - Returns: - y(tensor): batch x height x width - cb(tensor): batch x height/2 x width/2 - cr(tensor): batch x height/2 x width/2 - """ - image_2 = image.permute(0, 3, 1, 2).clone() - cb = F.avg_pool2d(image_2[:, 1, :, :].unsqueeze(1), kernel_size=2, stride=(2, 2), count_include_pad=False) - cr = F.avg_pool2d(image_2[:, 2, :, :].unsqueeze(1), kernel_size=2, stride=(2, 2), count_include_pad=False) - cb = cb.permute(0, 2, 3, 1) - cr = cr.permute(0, 2, 3, 1) - return image[:, :, :, 0], cb.squeeze(3), cr.squeeze(3) - - -class BlockSplitting(nn.Module): - """ Splitting image into patches - """ - - def __init__(self): - super(BlockSplitting, self).__init__() - self.k = 8 - - def forward(self, image): - """ - Args: - image(tensor): batch x height x width - - Returns: - Tensor: batch x h*w/64 x h x w - """ - height, _ = image.shape[1:3] - batch_size = image.shape[0] - image_reshaped = image.view(batch_size, height // self.k, self.k, -1, self.k) - image_transposed = image_reshaped.permute(0, 1, 3, 2, 4) - return image_transposed.contiguous().view(batch_size, -1, self.k, self.k) - - -class DCT8x8(nn.Module): - """ Discrete Cosine Transformation - """ - - def __init__(self): - super(DCT8x8, self).__init__() - tensor = np.zeros((8, 8, 8, 8), dtype=np.float32) - for x, y, u, v in itertools.product(range(8), repeat=4): - tensor[x, y, u, v] = np.cos((2 * x + 1) * u * np.pi / 16) * np.cos((2 * y + 1) * v * np.pi / 16) - alpha = np.array([1. / np.sqrt(2)] + [1] * 7) - self.tensor = nn.Parameter(torch.from_numpy(tensor).float()) - self.scale = nn.Parameter(torch.from_numpy(np.outer(alpha, alpha) * 0.25).float()) - - def forward(self, image): - """ - Args: - image(tensor): batch x height x width - - Returns: - Tensor: batch x height x width - """ - image = image - 128 - result = self.scale * torch.tensordot(image, self.tensor, dims=2) - result.view(image.shape) - return result - - -class YQuantize(nn.Module): - """ JPEG Quantization for Y channel - - Args: - rounding(function): rounding function to use - """ - - def __init__(self, rounding): - super(YQuantize, self).__init__() - self.rounding = rounding - self.y_table = y_table - - def forward(self, image, factor=1): - """ - Args: - image(tensor): batch x height x width - - Returns: - Tensor: batch x height x width - """ - if isinstance(factor, (int, float)): - image = image.float() / (self.y_table * factor) - else: - b = factor.size(0) - table = self.y_table.expand(b, 1, 8, 8) * factor.view(b, 1, 1, 1) - image = image.float() / table - image = self.rounding(image) - return image - - -class CQuantize(nn.Module): - """ JPEG Quantization for CbCr channels - - Args: - rounding(function): rounding function to use - """ - - def __init__(self, rounding): - super(CQuantize, self).__init__() - self.rounding = rounding - self.c_table = c_table - - def forward(self, image, factor=1): - """ - Args: - image(tensor): batch x height x width - - Returns: - Tensor: batch x height x width - """ - if isinstance(factor, (int, float)): - image = image.float() / (self.c_table * factor) - else: - b = factor.size(0) - table = self.c_table.expand(b, 1, 8, 8) * factor.view(b, 1, 1, 1) - image = image.float() / table - image = self.rounding(image) - return image - - -class CompressJpeg(nn.Module): - """Full JPEG compression algorithm - - Args: - rounding(function): rounding function to use - """ - - def __init__(self, rounding=torch.round): - super(CompressJpeg, self).__init__() - self.l1 = nn.Sequential(RGB2YCbCrJpeg(), ChromaSubsampling()) - self.l2 = nn.Sequential(BlockSplitting(), DCT8x8()) - self.c_quantize = CQuantize(rounding=rounding) - self.y_quantize = YQuantize(rounding=rounding) - - def forward(self, image, factor=1): - """ - Args: - image(tensor): batch x 3 x height x width - - Returns: - dict(tensor): Compressed tensor with batch x h*w/64 x 8 x 8. - """ - y, cb, cr = self.l1(image * 255) - components = {'y': y, 'cb': cb, 'cr': cr} - for k in components.keys(): - comp = self.l2(components[k]) - if k in ('cb', 'cr'): - comp = self.c_quantize(comp, factor=factor) - else: - comp = self.y_quantize(comp, factor=factor) - - components[k] = comp - - return components['y'], components['cb'], components['cr'] - - -# ------------------------ decompression ------------------------# - - -class YDequantize(nn.Module): - """Dequantize Y channel - """ - - def __init__(self): - super(YDequantize, self).__init__() - self.y_table = y_table - - def forward(self, image, factor=1): - """ - Args: - image(tensor): batch x height x width - - Returns: - Tensor: batch x height x width - """ - if isinstance(factor, (int, float)): - out = image * (self.y_table * factor) - else: - b = factor.size(0) - table = self.y_table.expand(b, 1, 8, 8) * factor.view(b, 1, 1, 1) - out = image * table - return out - - -class CDequantize(nn.Module): - """Dequantize CbCr channel - """ - - def __init__(self): - super(CDequantize, self).__init__() - self.c_table = c_table - - def forward(self, image, factor=1): - """ - Args: - image(tensor): batch x height x width - - Returns: - Tensor: batch x height x width - """ - if isinstance(factor, (int, float)): - out = image * (self.c_table * factor) - else: - b = factor.size(0) - table = self.c_table.expand(b, 1, 8, 8) * factor.view(b, 1, 1, 1) - out = image * table - return out - - -class iDCT8x8(nn.Module): - """Inverse discrete Cosine Transformation - """ - - def __init__(self): - super(iDCT8x8, self).__init__() - alpha = np.array([1. / np.sqrt(2)] + [1] * 7) - self.alpha = nn.Parameter(torch.from_numpy(np.outer(alpha, alpha)).float()) - tensor = np.zeros((8, 8, 8, 8), dtype=np.float32) - for x, y, u, v in itertools.product(range(8), repeat=4): - tensor[x, y, u, v] = np.cos((2 * u + 1) * x * np.pi / 16) * np.cos((2 * v + 1) * y * np.pi / 16) - self.tensor = nn.Parameter(torch.from_numpy(tensor).float()) - - def forward(self, image): - """ - Args: - image(tensor): batch x height x width - - Returns: - Tensor: batch x height x width - """ - image = image * self.alpha - result = 0.25 * torch.tensordot(image, self.tensor, dims=2) + 128 - result.view(image.shape) - return result - - -class BlockMerging(nn.Module): - """Merge patches into image - """ - - def __init__(self): - super(BlockMerging, self).__init__() - - def forward(self, patches, height, width): - """ - Args: - patches(tensor) batch x height*width/64, height x width - height(int) - width(int) - - Returns: - Tensor: batch x height x width - """ - k = 8 - batch_size = patches.shape[0] - image_reshaped = patches.view(batch_size, height // k, width // k, k, k) - image_transposed = image_reshaped.permute(0, 1, 3, 2, 4) - return image_transposed.contiguous().view(batch_size, height, width) - - -class ChromaUpsampling(nn.Module): - """Upsample chroma layers - """ - - def __init__(self): - super(ChromaUpsampling, self).__init__() - - def forward(self, y, cb, cr): - """ - Args: - y(tensor): y channel image - cb(tensor): cb channel - cr(tensor): cr channel - - Returns: - Tensor: batch x height x width x 3 - """ - - def repeat(x, k=2): - height, width = x.shape[1:3] - x = x.unsqueeze(-1) - x = x.repeat(1, 1, k, k) - x = x.view(-1, height * k, width * k) - return x - - cb = repeat(cb) - cr = repeat(cr) - return torch.cat([y.unsqueeze(3), cb.unsqueeze(3), cr.unsqueeze(3)], dim=3) - - -class YCbCr2RGBJpeg(nn.Module): - """Converts YCbCr image to RGB JPEG - """ - - def __init__(self): - super(YCbCr2RGBJpeg, self).__init__() - - matrix = np.array([[1., 0., 1.402], [1, -0.344136, -0.714136], [1, 1.772, 0]], dtype=np.float32).T - self.shift = nn.Parameter(torch.tensor([0, -128., -128.])) - self.matrix = nn.Parameter(torch.from_numpy(matrix)) - - def forward(self, image): - """ - Args: - image(tensor): batch x height x width x 3 - - Returns: - Tensor: batch x 3 x height x width - """ - result = torch.tensordot(image + self.shift, self.matrix, dims=1) - return result.view(image.shape).permute(0, 3, 1, 2) - - -class DeCompressJpeg(nn.Module): - """Full JPEG decompression algorithm - - Args: - rounding(function): rounding function to use - """ - - def __init__(self, rounding=torch.round): - super(DeCompressJpeg, self).__init__() - self.c_dequantize = CDequantize() - self.y_dequantize = YDequantize() - self.idct = iDCT8x8() - self.merging = BlockMerging() - self.chroma = ChromaUpsampling() - self.colors = YCbCr2RGBJpeg() - - def forward(self, y, cb, cr, imgh, imgw, factor=1): - """ - Args: - compressed(dict(tensor)): batch x h*w/64 x 8 x 8 - imgh(int) - imgw(int) - factor(float) - - Returns: - Tensor: batch x 3 x height x width - """ - components = {'y': y, 'cb': cb, 'cr': cr} - for k in components.keys(): - if k in ('cb', 'cr'): - comp = self.c_dequantize(components[k], factor=factor) - height, width = int(imgh / 2), int(imgw / 2) - else: - comp = self.y_dequantize(components[k], factor=factor) - height, width = imgh, imgw - comp = self.idct(comp) - components[k] = self.merging(comp, height, width) - # - image = self.chroma(components['y'], components['cb'], components['cr']) - image = self.colors(image) - - image = torch.min(255 * torch.ones_like(image), torch.max(torch.zeros_like(image), image)) - return image / 255 - - -# ------------------------ main DiffJPEG ------------------------ # - - -class DiffJPEG(nn.Module): - """This JPEG algorithm result is slightly different from cv2. - DiffJPEG supports batch processing. - - Args: - differentiable(bool): If True, uses custom differentiable rounding function, if False, uses standard torch.round - """ - - def __init__(self, differentiable=True): - super(DiffJPEG, self).__init__() - if differentiable: - rounding = diff_round - else: - rounding = torch.round - - self.compress = CompressJpeg(rounding=rounding) - self.decompress = DeCompressJpeg(rounding=rounding) - - def forward(self, x, quality): - """ - Args: - x (Tensor): Input image, bchw, rgb, [0, 1] - quality(float): Quality factor for jpeg compression scheme. - """ - factor = quality - if isinstance(factor, (int, float)): - factor = quality_to_factor(factor) - else: - for i in range(factor.size(0)): - factor[i] = quality_to_factor(factor[i]) - h, w = x.size()[-2:] - h_pad, w_pad = 0, 0 - # why should use 16 - if h % 16 != 0: - h_pad = 16 - h % 16 - if w % 16 != 0: - w_pad = 16 - w % 16 - x = F.pad(x, (0, w_pad, 0, h_pad), mode='constant', value=0) - - y, cb, cr = self.compress(x, factor=factor) - recovered = self.decompress(y, cb, cr, (h + h_pad), (w + w_pad), factor=factor) - recovered = recovered[:, :, 0:h, 0:w] - return recovered diff --git a/utils/image/usm_sharp.py b/utils/image/usm_sharp.py deleted file mode 100644 index 7b83532..0000000 --- a/utils/image/usm_sharp.py +++ /dev/null @@ -1,29 +0,0 @@ -# https://github.com/XPixelGroup/BasicSR/blob/master/basicsr/utils/img_process_util.py -import cv2 -import numpy as np -import torch - -from .common import filter2D - - -class USMSharp(torch.nn.Module): - - def __init__(self, radius=50, sigma=0): - super(USMSharp, self).__init__() - if radius % 2 == 0: - radius += 1 - self.radius = radius - kernel = cv2.getGaussianKernel(radius, sigma) - kernel = torch.FloatTensor(np.dot(kernel, kernel.transpose())).unsqueeze_(0) - self.register_buffer('kernel', kernel) - - def forward(self, img, weight=0.5, threshold=10): - blur = filter2D(img, self.kernel) - residual = img - blur - - mask = torch.abs(residual) * 255 > threshold - mask = mask.float() - soft_mask = filter2D(mask, self.kernel) - sharp = img + weight * residual - sharp = torch.clip(sharp, 0, 1) - return soft_mask * sharp + (1 - soft_mask) * img