Remove Numpy and CV2 to fully use Tensors

embedding_converter/src/dataset.py

@@ -1,9 +1,8 @@
 import glob
 
-import cv2
 from torch import Tensor
 from torch.utils.data import Dataset
-from torchvision import transforms
+from torchvision import transforms, io
 
 from .types import Batch
 
@@ -15,8 +14,8 @@ class StaticDataset(Dataset[Tensor]):
 
 	def __getitem__(self, index : int) -> Batch:
 		file_path = self.file_paths[index]
-		vision_frame = cv2.imread(file_path)
-		return self.transforms(vision_frame)
+		temp_tensor = io.read_image(file_path)
+		return self.transforms(temp_tensor)
 
 	def __len__(self) -> int:
 		return len(self.file_paths)
@@ -29,6 +28,5 @@ class StaticDataset(Dataset[Tensor]):
 			transforms.Resize((112, 112), interpolation = transforms.InterpolationMode.BICUBIC),
 			transforms.ColorJitter(brightness = 0.2, contrast = 0.2, saturation = 0.2, hue = 0.1),
 			transforms.ToTensor(),
-			transforms.Lambda(lambda temp_tensor: temp_tensor[[2, 1, 0], :, :]),
-			transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
+			transforms.Normalize(0.5, 0.5)
 		])

embedding_converter/src/types.py

@@ -1,10 +1,8 @@
 from typing import Any, TypeAlias
 
-from numpy.typing import NDArray
 from torch import Tensor
 
 Batch : TypeAlias = Tensor
 Embedding : TypeAlias = Tensor
-VisionFrame : TypeAlias = NDArray[Any]
 
 OptimizerConfig : TypeAlias = Any