WIP: Add custom dataset class to load single images and videos

Author: Turkka Helinoja

requirements.txt

@@ -9,3 +9,4 @@ json_tricks
 scikit-image
 yacs>=0.1.5
 tensorboardX
+Pillow

tools/custom_dataset.py

@@ -0,0 +1,69 @@
+import queue
+import pathlib
+
+from torch.utils.data import Dataset, IterableDataset
+import torchvision.transforms
+import torchvision.transforms.functional as TF
+
+from PIL import Image
+
+class CustomDataset(Dataset):
+    """Custom Dataset for handling single images and videos.
+
+    Custom torch.utils.data.Dataset class for handling single images and videos
+    and creating Dataset objects from them.
+    """
+    def __init__(self, path: str, transform: torchvision.transforms = None) -> None:
+        """Initializes CustomDataset
+
+        Args:
+            path (str): Path to the image or video
+            transform (torchvision.transforms, optional):
+                Transfrom that's used for transforming the images to tensors.
+                Defaults to None.
+        """
+        super().__init__()
+        self.path = path
+        self.transform = transform
+        self.queue = self.__load_data(self.path)
+
+
+    def __load_data(self, path: str) -> queue.Queue:
+        """Loads the data from the path. Data can be image or video.
+        TODO: Implement video support.
+
+        Args:
+            path (str): Path to the image or video
+
+        Returns:
+            queue.Queue: Queue which has the images loaded
+        """
+        buffer = queue.Queue()
+        file_extension = pathlib.Path(path).suffix
+        if file_extension == ".jpg" or file_extension == ".png":
+            image = Image.open(path)
+            #buffer.put(TF.to_tensor(image))
+            buffer.put(image)
+        return buffer
+
+    def __read_next_image(self):
+        while self.queue.qsize() > 0:
+            image = self.queue.get()
+            if self.transform is not None:
+                image = self.transform(image)
+            yield image
+
+        return None
+
+    def __iter__(self):
+        return self.__read_next_image()
+
+    def __len__(self):
+        return self.queue.qsize()
+
+    def __getitem__(self, x):
+        image = self.queue.get()
+        if self.transform is not None:
+            image = self.transform(image)
+
+        return image, None

tools/run_and_visualize.py

@@ -21,6 +21,8 @@
 from core.function import validate
 from utils.utils import create_logger
 
+from custom_dataset import CustomDataset
+
 import dataset
 import models
 
@@ -113,11 +115,13 @@ def main():
 
     # Load data
     if args.input != "":
-        with open(args.input, "r") as image:
             # TODO: Write a way to handle single images
             # TODO: Handle visualization
-            pass
-    elif args.video:
+            valid_dataset = CustomDataset(args.input, transforms.Compose([
+                transforms.ToTensor(),
+                normalize,
+            ]))
+    elif args.video != "":
         # TODO: Write a way to handle videos image by image
         # TODO: Handle visualization
         pass
@@ -131,17 +135,17 @@ def main():
             ])
         )
 
-        valid_loader = torch.utils.data.DataLoader(
-            valid_dataset,
-            batch_size=cfg.TEST.BATCH_SIZE_PER_GPU*len(cfg.GPUS),
-            shuffle=False,
-            num_workers=cfg.WORKERS,
-            pin_memory=True
-        )
+    valid_loader = torch.utils.data.DataLoader(
+        valid_dataset,
+        batch_size=cfg.TEST.BATCH_SIZE_PER_GPU*len(cfg.GPUS),
+        shuffle=False,
+        num_workers=cfg.WORKERS,
+        pin_memory=True
+    )
 
-        # evaluate on validation set
-        validate(cfg, valid_loader, valid_dataset, model, criterion,
-             final_output_dir, tb_log_dir)
+    # evaluate on validation set
+    validate(cfg, valid_loader, valid_dataset, model, criterion,
+            final_output_dir, tb_log_dir)
 
 if __name__ == '__main__':
     main()