Modernize data loader, remove read image helper

embedding_converter/src/data_loader.py

@@ -1,34 +1,29 @@
 import glob
 import random
 
+import cv2
 import torch
-import torchvision.transforms as transforms
+import transforms from torchvision
 from torch.utils.data import Dataset
 
-from .helper import read_image
 from .types import Batch, Paths
 
 
 class DataLoaderRecognition(Dataset[torch.Tensor]):
 	def __init__(self, dataset_file_pattern : str) -> None:
-		self.image_paths = self.prepare_image_paths(dataset_file_pattern)
-		self.dataset_total = len(self.image_paths)
+		self.image_paths = glob.glob(dataset_file_pattern)
 		self.transforms = self.compose_transforms()
 
 	def __getitem__(self, index : int) -> Batch:
-		target_image_path = random.choice(self.image_paths)
-		target_vision_frame = read_image(target_image_path)
-		target_tensor = self.transforms(target_vision_frame)
-		return target_tensor
+		image_path = random.choice(self.image_paths)
+		vision_frame = cv2.imread(image_path)
+		return self.transforms(vision_frame)
 
 	def __len__(self) -> int:
-		return self.dataset_total
-
-	def prepare_image_paths(self, dataset_file_pattern : str) -> Paths:
-		return glob.glob(dataset_file_pattern)
+		return len(self.image_paths)
 
 	def compose_transforms(self) -> transforms:
-		transform = transforms.Compose(
+		return transforms.Compose(
 		[
 			transforms.ToPILImage(),
 			transforms.Resize((112, 112), interpolation = transforms.InterpolationMode.BICUBIC),
@@ -37,4 +32,3 @@ class DataLoaderRecognition(Dataset[torch.Tensor]):
 			transforms.Lambda(lambda temp_tensor : temp_tensor[[2, 1, 0], :, :]),
 			transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
 		])
-		return transform

embedding_converter/src/helper.py

@@ -1,7 +0,0 @@
-import cv2
-
-from .types import VisionFrame
-
-
-def read_image(image_path : str) -> VisionFrame:
-	return cv2.imread(image_path)