clean data_loader

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

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]