From 95ea0a592ffccadc9986447e2ddfe18af3c9a9d3 Mon Sep 17 00:00:00 2001 From: James Betker Date: Wed, 16 Mar 2022 12:05:56 -0600 Subject: [PATCH] More cleaning --- codes/data/__init__.py | 24 +++++++++---------- codes/data/images/__init__.py | 0 .../base_unsupervised_image_dataset.py | 4 ++-- codes/data/{ => images}/byol_attachment.py | 0 .../data/{ => images}/chunk_with_reference.py | 0 codes/data/{ => images}/cifar.py | 0 codes/data/{ => images}/full_image_dataset.py | 0 codes/data/{ => images}/image_corruptor.py | 0 .../data/{ => images}/image_folder_dataset.py | 11 +++------ codes/data/{ => images}/image_label_parser.py | 0 ..._pair_with_corresponding_points_dataset.py | 16 ++----------- .../data/{ => images}/multi_frame_dataset.py | 2 +- codes/data/{ => images}/multiscale_dataset.py | 6 +---- .../data/{ => images}/paired_frame_dataset.py | 2 +- codes/data/{ => images}/random_dataset.py | 0 .../data/{ => images}/single_image_dataset.py | 3 +-- codes/data/{ => images}/stylegan2_dataset.py | 0 codes/data/{ => images}/zip_file_dataset.py | 0 codes/data/torch_dataset.py | 3 +-- .../image_latents/byol/byol_model_wrapper.py | 2 +- .../image_latents/byol/byol_structural.py | 2 +- codes/scripts/byol/byol_resnet_playground.py | 2 +- .../scripts/byol/byol_segformer_playground.py | 2 +- .../scripts/byol/byol_spinenet_playground.py | 2 +- .../scripts/srflow_latent_space_playground.py | 2 +- .../ui/image_labeler/image_labeler_ui.py | 2 +- .../scripts/ui/image_labeler/label_editor.py | 2 +- codes/scripts/use_generator_as_filter.py | 2 +- .../progressive_zoom.py | 2 +- codes/trainer/eval/flow_gaussian_nll.py | 2 +- .../single_point_pair_contrastive_eval.py | 9 ++----- codes/trainer/eval/sr_style.py | 2 +- 32 files changed, 38 insertions(+), 66 deletions(-) create mode 100644 codes/data/images/__init__.py rename codes/data/{ => images}/base_unsupervised_image_dataset.py (98%) rename codes/data/{ => images}/byol_attachment.py (100%) rename codes/data/{ => images}/chunk_with_reference.py (100%) rename codes/data/{ => images}/cifar.py (100%) rename codes/data/{ => images}/full_image_dataset.py (100%) rename codes/data/{ => images}/image_corruptor.py (100%) rename codes/data/{ => images}/image_folder_dataset.py (98%) rename codes/data/{ => images}/image_label_parser.py (100%) rename codes/data/{ => images}/image_pair_with_corresponding_points_dataset.py (89%) rename codes/data/{ => images}/multi_frame_dataset.py (98%) rename codes/data/{ => images}/multiscale_dataset.py (97%) rename codes/data/{ => images}/paired_frame_dataset.py (97%) rename codes/data/{ => images}/random_dataset.py (100%) rename codes/data/{ => images}/single_image_dataset.py (96%) rename codes/data/{ => images}/stylegan2_dataset.py (100%) rename codes/data/{ => images}/zip_file_dataset.py (100%) diff --git a/codes/data/__init__.py b/codes/data/__init__.py index 97246259..a390d4ea 100644 --- a/codes/data/__init__.py +++ b/codes/data/__init__.py @@ -33,33 +33,33 @@ def create_dataset(dataset_opt, return_collate=False): # datasets for image restoration if mode == 'fullimage': - from data.full_image_dataset import FullImageDataset as D + from data.images.full_image_dataset import FullImageDataset as D elif mode == 'single_image_extensible': - from data.single_image_dataset import SingleImageDataset as D + from data.images.single_image_dataset import SingleImageDataset as D elif mode == 'multi_frame_extensible': - from data.multi_frame_dataset import MultiFrameDataset as D + from data.images.multi_frame_dataset import MultiFrameDataset as D elif mode == 'combined': from data.combined_dataset import CombinedDataset as D elif mode == 'multiscale': - from data.multiscale_dataset import MultiScaleDataset as D + from data.images.multiscale_dataset import MultiScaleDataset as D elif mode == 'paired_frame': - from data.paired_frame_dataset import PairedFrameDataset as D + from data.images.paired_frame_dataset import PairedFrameDataset as D elif mode == 'stylegan2': - from data.stylegan2_dataset import Stylegan2Dataset as D + from data.images.stylegan2_dataset import Stylegan2Dataset as D elif mode == 'imagefolder': - from data.image_folder_dataset import ImageFolderDataset as D + from data.images.image_folder_dataset import ImageFolderDataset as D elif mode == 'torch_dataset': from data.torch_dataset import TorchDataset as D elif mode == 'byol_dataset': - from data.byol_attachment import ByolDatasetWrapper as D + from data.images.byol_attachment import ByolDatasetWrapper as D elif mode == 'byol_structured_dataset': - from data.byol_attachment import StructuredCropDatasetWrapper as D + from data.images.byol_attachment import StructuredCropDatasetWrapper as D elif mode == 'random_aug_wrapper': - from data.byol_attachment import DatasetRandomAugWrapper as D + from data.images.byol_attachment import DatasetRandomAugWrapper as D elif mode == 'random_dataset': - from data.random_dataset import RandomDataset as D + from data.images.random_dataset import RandomDataset as D elif mode == 'zipfile': - from data.zip_file_dataset import ZipFileDataset as D + from data.images.zip_file_dataset import ZipFileDataset as D elif mode == 'nv_tacotron': from data.audio.nv_tacotron_dataset import TextWavLoader as D from data.audio.nv_tacotron_dataset import TextMelCollate as C diff --git a/codes/data/images/__init__.py b/codes/data/images/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/codes/data/base_unsupervised_image_dataset.py b/codes/data/images/base_unsupervised_image_dataset.py similarity index 98% rename from codes/data/base_unsupervised_image_dataset.py rename to codes/data/images/base_unsupervised_image_dataset.py index 83ae0a3b..ce8b8c3f 100644 --- a/codes/data/base_unsupervised_image_dataset.py +++ b/codes/data/images/base_unsupervised_image_dataset.py @@ -1,7 +1,7 @@ import torch from torch.utils import data -from data.image_corruptor import ImageCorruptor -from data.chunk_with_reference import ChunkWithReference +from data.images.image_corruptor import ImageCorruptor +from data.images.chunk_with_reference import ChunkWithReference import os import cv2 import numpy as np diff --git a/codes/data/byol_attachment.py b/codes/data/images/byol_attachment.py similarity index 100% rename from codes/data/byol_attachment.py rename to codes/data/images/byol_attachment.py diff --git a/codes/data/chunk_with_reference.py b/codes/data/images/chunk_with_reference.py similarity index 100% rename from codes/data/chunk_with_reference.py rename to codes/data/images/chunk_with_reference.py diff --git a/codes/data/cifar.py b/codes/data/images/cifar.py similarity index 100% rename from codes/data/cifar.py rename to codes/data/images/cifar.py diff --git a/codes/data/full_image_dataset.py b/codes/data/images/full_image_dataset.py similarity index 100% rename from codes/data/full_image_dataset.py rename to codes/data/images/full_image_dataset.py diff --git a/codes/data/image_corruptor.py b/codes/data/images/image_corruptor.py similarity index 100% rename from codes/data/image_corruptor.py rename to codes/data/images/image_corruptor.py diff --git a/codes/data/image_folder_dataset.py b/codes/data/images/image_folder_dataset.py similarity index 98% rename from codes/data/image_folder_dataset.py rename to codes/data/images/image_folder_dataset.py index 712d04dd..e046eed1 100644 --- a/codes/data/image_folder_dataset.py +++ b/codes/data/images/image_folder_dataset.py @@ -1,24 +1,19 @@ import functools -import glob -import itertools import random import cv2 -import kornia import numpy as np -import pytorch_ssim import torch -import os import torchvision from torch.utils.data import DataLoader -from torchvision.transforms import Normalize, CenterCrop +from torchvision.transforms import Normalize from tqdm import tqdm from data import util # Builds a dataset created from a simple folder containing a list of training/test/validation images. -from data.image_corruptor import ImageCorruptor, kornia_color_jitter_numpy -from data.image_label_parser import VsNetImageLabeler +from data.images.image_corruptor import ImageCorruptor, kornia_color_jitter_numpy +from data.images.image_label_parser import VsNetImageLabeler from utils.util import opt_get diff --git a/codes/data/image_label_parser.py b/codes/data/images/image_label_parser.py similarity index 100% rename from codes/data/image_label_parser.py rename to codes/data/images/image_label_parser.py diff --git a/codes/data/image_pair_with_corresponding_points_dataset.py b/codes/data/images/image_pair_with_corresponding_points_dataset.py similarity index 89% rename from codes/data/image_pair_with_corresponding_points_dataset.py rename to codes/data/images/image_pair_with_corresponding_points_dataset.py index c6650080..1d2d35c0 100644 --- a/codes/data/image_pair_with_corresponding_points_dataset.py +++ b/codes/data/images/image_pair_with_corresponding_points_dataset.py @@ -1,11 +1,3 @@ -import glob -import itertools -import random - -import cv2 -import kornia -import numpy as np -import pytorch_ssim import torch import os @@ -13,14 +5,10 @@ import torchvision from PIL import Image from torch.utils.data import DataLoader, Dataset from torchvision import transforms -from torchvision.transforms import Normalize from tqdm import tqdm -from data import util + # Builds a dataset created from a simple folder containing a list of training/test/validation images. -from data.image_corruptor import ImageCorruptor -from data.image_label_parser import VsNetImageLabeler -from utils.util import opt_get class ImagePairWithCorrespondingPointsDataset(Dataset): @@ -70,7 +58,7 @@ if __name__ == '__main__': 'path': 'F:\\dlas\\codes\\scripts\\ui\\image_pair_labeler\\results', 'size': 256 } - output_path = '.' + output_path = '..' ds = DataLoader(ImagePairWithCorrespondingPointsDataset(opt), shuffle=True, num_workers=0) for i, d in tqdm(enumerate(ds)): diff --git a/codes/data/multi_frame_dataset.py b/codes/data/images/multi_frame_dataset.py similarity index 98% rename from codes/data/multi_frame_dataset.py rename to codes/data/images/multi_frame_dataset.py index ce99b348..76c4b127 100644 --- a/codes/data/multi_frame_dataset.py +++ b/codes/data/images/multi_frame_dataset.py @@ -1,4 +1,4 @@ -from data.base_unsupervised_image_dataset import BaseUnsupervisedImageDataset +from data.images.base_unsupervised_image_dataset import BaseUnsupervisedImageDataset import numpy as np import torch from bisect import bisect_left diff --git a/codes/data/multiscale_dataset.py b/codes/data/images/multiscale_dataset.py similarity index 97% rename from codes/data/multiscale_dataset.py rename to codes/data/images/multiscale_dataset.py index 2314e3e6..540d3076 100644 --- a/codes/data/multiscale_dataset.py +++ b/codes/data/images/multiscale_dataset.py @@ -4,13 +4,9 @@ import cv2 import torch import torch.utils.data as data import data.util as util -from PIL import Image, ImageOps -from io import BytesIO -import torchvision.transforms.functional as F - # Reads full-quality images and pulls tiles at regular zoom intervals from them. Only usable for training purposes. -from data.image_corruptor import ImageCorruptor +from data.images.image_corruptor import ImageCorruptor # Selects the smallest dimension from the image and crops it randomly so the other dimension matches. The cropping diff --git a/codes/data/paired_frame_dataset.py b/codes/data/images/paired_frame_dataset.py similarity index 97% rename from codes/data/paired_frame_dataset.py rename to codes/data/images/paired_frame_dataset.py index 49328abc..09605995 100644 --- a/codes/data/paired_frame_dataset.py +++ b/codes/data/images/paired_frame_dataset.py @@ -1,4 +1,4 @@ -from data.base_unsupervised_image_dataset import BaseUnsupervisedImageDataset +from data.images.base_unsupervised_image_dataset import BaseUnsupervisedImageDataset import numpy as np import torch from bisect import bisect_left diff --git a/codes/data/random_dataset.py b/codes/data/images/random_dataset.py similarity index 100% rename from codes/data/random_dataset.py rename to codes/data/images/random_dataset.py diff --git a/codes/data/single_image_dataset.py b/codes/data/images/single_image_dataset.py similarity index 96% rename from codes/data/single_image_dataset.py rename to codes/data/images/single_image_dataset.py index 4048f197..3aaebb56 100644 --- a/codes/data/single_image_dataset.py +++ b/codes/data/images/single_image_dataset.py @@ -2,8 +2,7 @@ import random from bisect import bisect_left import numpy as np import torch -from torch.utils import data -from data.base_unsupervised_image_dataset import BaseUnsupervisedImageDataset +from data.images.base_unsupervised_image_dataset import BaseUnsupervisedImageDataset # Builds a dataset composed of a set of folders. Each folder represents a single high resolution image that has been diff --git a/codes/data/stylegan2_dataset.py b/codes/data/images/stylegan2_dataset.py similarity index 100% rename from codes/data/stylegan2_dataset.py rename to codes/data/images/stylegan2_dataset.py diff --git a/codes/data/zip_file_dataset.py b/codes/data/images/zip_file_dataset.py similarity index 100% rename from codes/data/zip_file_dataset.py rename to codes/data/images/zip_file_dataset.py diff --git a/codes/data/torch_dataset.py b/codes/data/torch_dataset.py index 079cf4fe..7ad90650 100644 --- a/codes/data/torch_dataset.py +++ b/codes/data/torch_dataset.py @@ -1,10 +1,9 @@ -import torch from torch.utils.data import Dataset import torchvision.transforms as T from torchvision import datasets # Wrapper for basic pytorch datasets which re-wraps them into a format usable by ExtensibleTrainer. -from data.cifar import CIFAR100, CIFAR10 +from data.images.cifar import CIFAR100, CIFAR10 from utils.util import opt_get diff --git a/codes/models/image_latents/byol/byol_model_wrapper.py b/codes/models/image_latents/byol/byol_model_wrapper.py index 9c3b4f3f..cf3b6ea0 100644 --- a/codes/models/image_latents/byol/byol_model_wrapper.py +++ b/codes/models/image_latents/byol/byol_model_wrapper.py @@ -9,7 +9,7 @@ import torchvision from kornia import filters from torch import nn -from data.byol_attachment import RandomApply +from data.images.byol_attachment import RandomApply from trainer.networks import register_model, create_model from utils.util import checkpoint, opt_get diff --git a/codes/models/image_latents/byol/byol_structural.py b/codes/models/image_latents/byol/byol_structural.py index 778d8eeb..dbb8fe5f 100644 --- a/codes/models/image_latents/byol/byol_structural.py +++ b/codes/models/image_latents/byol/byol_structural.py @@ -4,7 +4,7 @@ import torch import torch.nn.functional as F from torch import nn -from data.byol_attachment import reconstructed_shared_regions +from data.images.byol_attachment import reconstructed_shared_regions from models.image_latents.byol.byol_model_wrapper import singleton, EMA, get_module_device, set_requires_grad, \ update_moving_average from trainer.networks import create_model, register_model diff --git a/codes/scripts/byol/byol_resnet_playground.py b/codes/scripts/byol/byol_resnet_playground.py index b5014567..87274189 100644 --- a/codes/scripts/byol/byol_resnet_playground.py +++ b/codes/scripts/byol/byol_resnet_playground.py @@ -9,7 +9,7 @@ from torch.utils.data import DataLoader from torchvision.transforms import ToTensor from tqdm import tqdm -from data.image_folder_dataset import ImageFolderDataset +from data.images.image_folder_dataset import ImageFolderDataset from models.classifiers.resnet_with_checkpointing import resnet50 # Computes the structural euclidean distance between [x,y]. "Structural" here means the [h,w] dimensions are preserved diff --git a/codes/scripts/byol/byol_segformer_playground.py b/codes/scripts/byol/byol_segformer_playground.py index 2fc9db51..eccc2a48 100644 --- a/codes/scripts/byol/byol_segformer_playground.py +++ b/codes/scripts/byol/byol_segformer_playground.py @@ -9,7 +9,7 @@ from torch.utils.data import DataLoader from torchvision.transforms import ToTensor, Normalize from tqdm import tqdm -from data.image_folder_dataset import ImageFolderDataset +from data.images.image_folder_dataset import ImageFolderDataset from models.segformer.segformer import Segformer # Computes the structural euclidean distance between [x,y]. "Structural" here means the [h,w] dimensions are preserved diff --git a/codes/scripts/byol/byol_spinenet_playground.py b/codes/scripts/byol/byol_spinenet_playground.py index d646ba57..8f2625c8 100644 --- a/codes/scripts/byol/byol_spinenet_playground.py +++ b/codes/scripts/byol/byol_spinenet_playground.py @@ -10,7 +10,7 @@ from torchvision.transforms import ToTensor, Resize from tqdm import tqdm import numpy as np -from data.image_folder_dataset import ImageFolderDataset +from data.images.image_folder_dataset import ImageFolderDataset from models.image_latents.spinenet_arch import SpineNet diff --git a/codes/scripts/srflow_latent_space_playground.py b/codes/scripts/srflow_latent_space_playground.py index 8cbce442..a91b8528 100644 --- a/codes/scripts/srflow_latent_space_playground.py +++ b/codes/scripts/srflow_latent_space_playground.py @@ -13,7 +13,7 @@ from tqdm import tqdm import utils.options as option import utils -from data.image_corruptor import ImageCorruptor +from data.images.image_corruptor import ImageCorruptor from trainer.ExtensibleTrainer import ExtensibleTrainer from utils import util diff --git a/codes/scripts/ui/image_labeler/image_labeler_ui.py b/codes/scripts/ui/image_labeler/image_labeler_ui.py index 1eb7d15d..bbe5921a 100644 --- a/codes/scripts/ui/image_labeler/image_labeler_ui.py +++ b/codes/scripts/ui/image_labeler/image_labeler_ui.py @@ -11,7 +11,7 @@ import torch import torchvision from PIL import ImageTk -from data.image_label_parser import VsNetImageLabeler +from data.images.image_label_parser import VsNetImageLabeler from scripts.ui.image_labeler.pretrained_image_patch_classifier import PretrainedImagePatchClassifier # Globals used to define state that event handlers might operate on. diff --git a/codes/scripts/ui/image_labeler/label_editor.py b/codes/scripts/ui/image_labeler/label_editor.py index b081f40b..19e30c0d 100644 --- a/codes/scripts/ui/image_labeler/label_editor.py +++ b/codes/scripts/ui/image_labeler/label_editor.py @@ -1,6 +1,6 @@ import orjson -from data.image_label_parser import VsNetImageLabeler +from data.images.image_label_parser import VsNetImageLabeler # Translates from the label JSON output of the VS.NET UI to something more compact and usable. diff --git a/codes/scripts/use_generator_as_filter.py b/codes/scripts/use_generator_as_filter.py index f08f7d44..c4116ed0 100644 --- a/codes/scripts/use_generator_as_filter.py +++ b/codes/scripts/use_generator_as_filter.py @@ -3,7 +3,7 @@ import shutil from torch.utils.data import DataLoader -from data.single_image_dataset import SingleImageDataset +from data.images.single_image_dataset import SingleImageDataset from tqdm import tqdm import torch diff --git a/codes/trainer/custom_training_components/progressive_zoom.py b/codes/trainer/custom_training_components/progressive_zoom.py index 63f0f980..038db678 100644 --- a/codes/trainer/custom_training_components/progressive_zoom.py +++ b/codes/trainer/custom_training_components/progressive_zoom.py @@ -5,7 +5,7 @@ import torch import torchvision from torch.cuda.amp import autocast -from data.multiscale_dataset import build_multiscale_patch_index_map +from data.images.multiscale_dataset import build_multiscale_patch_index_map from trainer.inject import Injector from trainer.losses import extract_params_from_state import os.path as osp diff --git a/codes/trainer/eval/flow_gaussian_nll.py b/codes/trainer/eval/flow_gaussian_nll.py index 237d40e8..42d72f36 100644 --- a/codes/trainer/eval/flow_gaussian_nll.py +++ b/codes/trainer/eval/flow_gaussian_nll.py @@ -5,7 +5,7 @@ from tqdm import tqdm import trainer.eval.evaluator as evaluator # Evaluate how close to true Gaussian a flow network predicts in a "normal" pass given a LQ/HQ image pair. -from data.image_folder_dataset import ImageFolderDataset +from data.images.image_folder_dataset import ImageFolderDataset from models.image_generation.srflow import GaussianDiag diff --git a/codes/trainer/eval/single_point_pair_contrastive_eval.py b/codes/trainer/eval/single_point_pair_contrastive_eval.py index d73ce6cd..2b16e509 100644 --- a/codes/trainer/eval/single_point_pair_contrastive_eval.py +++ b/codes/trainer/eval/single_point_pair_contrastive_eval.py @@ -1,18 +1,13 @@ -import os - import torch -import os.path as osp -import torchvision from torch.nn import MSELoss from torch.utils.data import DataLoader from tqdm import tqdm import trainer.eval.evaluator as evaluator -from pytorch_fid import fid_score -from data.image_pair_with_corresponding_points_dataset import ImagePairWithCorrespondingPointsDataset +from data.images.image_pair_with_corresponding_points_dataset import ImagePairWithCorrespondingPointsDataset from models.segformer.segformer import Segformer -from utils.util import opt_get + # Uses two datasets: a "similar" and "dissimilar" dataset, each of which contains pairs of images and similar/dissimilar # points in those datasets. Uses the provided network to compute a latent vector for both similar and dissimilar. diff --git a/codes/trainer/eval/sr_style.py b/codes/trainer/eval/sr_style.py index add84616..7b50d619 100644 --- a/codes/trainer/eval/sr_style.py +++ b/codes/trainer/eval/sr_style.py @@ -11,7 +11,7 @@ from pytorch_fid import fid_score # Evaluate that feeds a LR structure into the input, then calculates a FID score on the results added to # the interpolated LR structure. -from data.stylegan2_dataset import Stylegan2Dataset +from data.images.stylegan2_dataset import Stylegan2Dataset class SrStyleTransferEvaluator(evaluator.Evaluator):