From 798ba48a52790210154de197c57187bd9ca70b47 Mon Sep 17 00:00:00 2001 From: harisreedhar Date: Tue, 21 Jan 2025 21:43:42 +0530 Subject: [PATCH] clean data_loader --- face_swapper/src/data_loader.py | 77 ++++++++++++++------------------- face_swapper/src/typing.py | 3 ++ 2 files changed, 36 insertions(+), 44 deletions(-) diff --git a/face_swapper/src/data_loader.py b/face_swapper/src/data_loader.py index cceae06..d203054 100644 --- a/face_swapper/src/data_loader.py +++ b/face_swapper/src/data_loader.py @@ -1,81 +1,70 @@ import configparser import glob +import os.path import random import cv2 +import torch import torchvision.transforms as transforms -import tqdm -from PIL import Image from torch.utils.data import TensorDataset -from .typing import Batch +from .typing import Batch, VisionFrame CONFIG = configparser.ConfigParser() CONFIG.read('config.ini') -def read_image(image_path: str) -> Image.Image: +def read_image(image_path: str) -> VisionFrame: image = cv2.imread(image_path)[:, :, ::-1] - pil_image = Image.fromarray(image) # @todo like said, use the PIL transformator - return pil_image + return image class DataLoaderVGG(TensorDataset): def __init__(self, dataset_path : str) -> None: - self.same_person_probability = float(CONFIG.get('preparing.dataloader', 'same_person_probability')) # @todo use CONFIG.getfloat() - also config block at the top + self.same_person_probability = CONFIG.getfloat('preparing.dataloader', 'same_person_probability') self.image_paths = glob.glob('{}/*/*.*g'.format(dataset_path)) # @todo globs belong to the config self.folder_paths = glob.glob('{}/*'.format(dataset_path)) - self.image_path_dict = {} # @todo we are not using dict as suffix... this image_path_set? + self.image_path_set = {} self._current_index = 0 - for folder_path in tqdm.tqdm(self.folder_paths): + for folder_path in self.folder_paths: image_paths = glob.glob('{}/*'.format(folder_path)) # @todo not sure about alls this globs being used here :-) - self.image_path_dict[folder_path] = image_paths + self.image_path_set[folder_path] = image_paths self.dataset_total = len(self.image_paths) - self.transforms_basic = transforms.Compose( + self.transforms = transforms.Compose( [ - transforms.Resize((256, 256), interpolation=transforms.InterpolationMode.BICUBIC), - transforms.ToTensor(), - transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)) - ]) - self.transforms_moderate = transforms.Compose( - [ - transforms.Resize((256, 256), interpolation=transforms.InterpolationMode.BICUBIC), - transforms.ToTensor(), + transforms.ToPILImage(), + transforms.Resize((256, 256), interpolation = transforms.InterpolationMode.BICUBIC), transforms.ColorJitter(brightness = 0.2, contrast = 0.2, saturation = 0.2, hue = 0.1), transforms.RandomAffine(4, translate = (0.01, 0.01), scale = (0.98, 1.02), shear = (1, 1), fill = 0), - transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)) - ]) - self.transforms_complex = transforms.Compose( - [ - transforms.Resize((256, 256), interpolation=transforms.InterpolationMode.BICUBIC), transforms.ToTensor(), - transforms.RandomHorizontalFlip(p = 0.5), - transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation = 0.2, hue = 0.1), - transforms.RandomAffine(8, translate = (0.02, 0.02), scale = (0.98, 1.02), shear = (1, 1), fill = 0), - transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)), - transforms.Resize((256, 256), interpolation=transforms.InterpolationMode.BICUBIC), + transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)) ]) def __getitem__(self, item : int) -> Batch: source_image_path = self.image_paths[item] - source = read_image(source_image_path) - if random.random() > self.same_person_probability: # @todo if -> we_call_a_method_that_explains_what_we_do() - is_same_person = 0 - target_image_path = random.choice(self.image_paths) - target = read_image(target_image_path) - source_transform = self.transforms_moderate(source) - target_transform = self.transforms_complex(target) - else: # @todo else -> we_do_some_alternative_action() - in other words, move it to speaking methods :-) - is_same_person = 1 - source_folder_path = '/'.join(source_image_path.split('/')[:-1]) # @todo use os.path.join() - target_image_path = random.choice(self.image_path_dict[source_folder_path]) - target = read_image(target_image_path) - source_transform = self.transforms_basic(source) - target_transform = self.transforms_basic(target) + if random.random() > self.same_person_probability: + return self.prepare_same_person(source_image_path) + return self.prepare_different_person(source_image_path) - return source_transform, target_transform, is_same_person + def prepare_different_person(self, source_image_path : str) -> Batch: + is_same_person = torch.tensor(0) + target_image_path = random.choice(self.image_paths) + source_vision_frame = read_image(source_image_path) + target_vision_frame = read_image(target_image_path) + source_tensor = self.transforms(source_vision_frame) + target_tensor = self.transforms(target_vision_frame) + return source_tensor, target_tensor, is_same_person + + def prepare_same_person(self, source_image_path : str) -> Batch: + is_same_person = torch.tensor(1) + target_image_path = random.choice(self.image_path_set.get(os.path.dirname(source_image_path))) + source_vision_frame = read_image(source_image_path) + target_vision_frame = read_image(target_image_path) + source_tensor = self.transforms(source_vision_frame) + target_tensor = self.transforms(target_vision_frame) + return source_tensor, target_tensor, is_same_person def __len__(self) -> int: return self.dataset_total diff --git a/face_swapper/src/typing.py b/face_swapper/src/typing.py index da89355..d72df42 100644 --- a/face_swapper/src/typing.py +++ b/face_swapper/src/typing.py @@ -1,9 +1,11 @@ from collections import OrderedDict from typing import Any, Dict, List, Tuple +from numpy.typing import NDArray from torch import Tensor from torch.utils.data import DataLoader + Batch = Tuple[Any, Any, Any] Loader = DataLoader[Tuple[Tensor, ...]] TargetAttributes = Tuple[Tensor, ...] @@ -17,3 +19,4 @@ VisionTensor = Tensor LossTensor = Tensor GeneratorLossSet = Dict[str, LossTensor] DiscriminatorLossSet = Dict[str, LossTensor] +VisionFrame = NDArray[Any]