Change dataset for LSdir

datasets.py

@@ -4,6 +4,7 @@ import os
 
 import torch
 from torch.utils.data import Dataset
+from PIL import Image
 
 
 class Preprocessed_fastMRI(torch.utils.data.Dataset):
@@ -82,3 +83,26 @@ class Preprocessed_LIDCIDRI(torch.utils.data.Dataset):
             img = self.transform(img)
 
         img = img.unsqueeze(0)  # add channel dim
+        return img
+
+
+class LsdirMiniDataset(torch.utils.data.Dataset):
+    def __init__(
+	self,
+	root: str,
+	transform: Optional[Callable] = None,
+    ) -> None:
+        self.root = root
+        self.image_files = [f for f in os.listdir(self.root) if f.lower().endswith('.png')]
+        self.transform = transform
+
+    def __len__(self) -> int:
+        return len(self.image_files)
+
+    def __getitem__(self, idx):
+        img_path = os.path.join(self.root, self.image_files[idx])
+        img = Image.open(img_path)
+        if self.transform:
+            img = self.transform(img)
+
+        return img

evals.py

@@ -6,7 +6,7 @@ import torch
 from deepinv.physics.generator import MotionBlurGenerator, SigmaGenerator
 from torchvision import transforms
 
-from datasets import Preprocessed_fastMRI, Preprocessed_LIDCIDRI
+from datasets import Preprocessed_fastMRI, Preprocessed_LIDCIDRI, LsdirMiniDataset
 from utils import get_model
 
 DEFAULT_MODEL_PARAMS = {
@@ -491,9 +491,7 @@ class BaselineModel(torch.nn.Module):
 
 
 class EvalDataset(torch.utils.data.Dataset):
-    """
-    We expect that images are 480x480.
-    """
+    """"""
     all_datasets = ["Natural", "MRI", "CT"]
 
     def __init__(self, dataset_name: str, device_str: str = "cpu") -> None:
@@ -504,9 +502,8 @@ class EvalDataset(torch.utils.data.Dataset):
         if self.name == 'Natural':
             self.root = 'img_samples/LSDIR_samples'
             self.transform = transforms.Compose([transforms.ToTensor()])
-            self.dataset = dinv.datasets.LsdirHR(root=self.root,
-                                                 download=False,
-                                                 transform=self.transform)
+            self.dataset = LsdirMiniDataset(root=self.root,
+                                            transform=self.transform)
         elif self.name == 'MRI':
             self.root = 'img_samples/FastMRI_samples'
             self.transform = transforms.CenterCrop((640, 320))  # , pad_if_needed=True)