Fixed issues in mssim and requirements

Author: AntixK

dataset.py

@@ -11,39 +11,32 @@
 import zipfile
 
 
-# class MyDataset(Dataset):
-#     def __init__(self):
-#         pass
+# Add your custom dataset class here
+class MyDataset(Dataset):
+    def __init__(self):
+        pass
 
 
-#     def __len__(self):
-#         pass
+    def __len__(self):
+        pass
 
-#     def __getitem__(self, idx):
-#         pass
+    def __getitem__(self, idx):
+        pass
+
 
 class MyCelebA(CelebA):
     """
+    A work-around to address issues with pytorch's celebA dataset class.
+    
     Download and Extract
     URL : https://drive.google.com/file/d/1m8-EBPgi5MRubrm6iQjafK2QMHDBMSfJ/view?usp=sharing
     """
 
     def _check_integrity(self) -> bool:
         return True
-
-class Food101(Dataset):
-    """
-    URL : https://data.vision.ee.ethz.ch/cvl/datasets_extra/food-101/
-    """
-    def __init__(self):
-        pass
 
-    def __len__(self):
-        pass
-    
-    def __getitem__(self, idx):
-        pass
 
+
 class OxfordPets(Dataset):
     """
     URL = https://www.robots.ox.ac.uk/~vgg/data/pets/
@@ -53,7 +46,7 @@ def __init__(self,
                  split: str,
                  transform: Callable,
                 **kwargs):
-        self.data_dir = Path(data_path)        
+        self.data_dir = Path(data_path) / "OxfordPets"        
         self.transforms = transform
         imgs = sorted([f for f in self.data_dir.iterdir() if f.suffix == '.jpg'])
 
@@ -107,59 +100,58 @@ def __init__(
     def setup(self, stage: Optional[str] = None) -> None:
 #       =========================  OxfordPets Dataset  =========================
 
-#         train_transforms = transforms.Compose([transforms.RandomHorizontalFlip(),
-#                                               transforms.CenterCrop(self.patch_size),
-# #                                               transforms.Resize(self.patch_size),
-#                                               transforms.ToTensor(),
-#                                                 transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))])
-        
-#         val_transforms = transforms.Compose([transforms.RandomHorizontalFlip(),
-#                                             transforms.CenterCrop(self.patch_size),
-# #                                             transforms.Resize(self.patch_size),
-#                                             transforms.ToTensor(),
-#                                               transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))])
-
-#         self.train_dataset = OxfordPets(
-#             self.data_dir,
-#             split='train',
-#             transform=train_transforms,
-#         )
-        
-#         self.val_dataset = OxfordPets(
-#             self.data_dir,
-#             split='val',
-#             transform=val_transforms,
-#         )
-        
-#       =========================  CelebA Dataset  =========================
-    
         train_transforms = transforms.Compose([transforms.RandomHorizontalFlip(),
-                                              transforms.CenterCrop(148),
-                                              transforms.Resize(self.patch_size),
-                                              transforms.ToTensor(),])
+                                              transforms.CenterCrop(self.patch_size),
+#                                               transforms.Resize(self.patch_size),
+                                              transforms.ToTensor(),
+                                                transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))])
 
         val_transforms = transforms.Compose([transforms.RandomHorizontalFlip(),
-                                            transforms.CenterCrop(148),
-                                            transforms.Resize(self.patch_size),
-                                            transforms.ToTensor(),])
-        
-        self.train_dataset = MyCelebA(
+                                            transforms.CenterCrop(self.patch_size),
+#                                             transforms.Resize(self.patch_size),
+                                            transforms.ToTensor(),
+                                              transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))])
+
+        self.train_dataset = OxfordPets(
             self.data_dir,
             split='train',
             transform=train_transforms,
-            download=False,
         )
 
-        # Replace CelebA with your dataset
-        self.val_dataset = MyCelebA(
+        self.val_dataset = OxfordPets(
             self.data_dir,
-            split='test',
+            split='val',
             transform=val_transforms,
-            download=False,
         )
+        
+#       =========================  CelebA Dataset  =========================
+    
+#         train_transforms = transforms.Compose([transforms.RandomHorizontalFlip(),
+#                                               transforms.CenterCrop(148),
+#                                               transforms.Resize(self.patch_size),
+#                                               transforms.ToTensor(),])
+        
+#         val_transforms = transforms.Compose([transforms.RandomHorizontalFlip(),
+#                                             transforms.CenterCrop(148),
+#                                             transforms.Resize(self.patch_size),
+#                                             transforms.ToTensor(),])
+        
+#         self.train_dataset = MyCelebA(
+#             self.data_dir,
+#             split='train',
+#             transform=train_transforms,
+#             download=False,
+#         )
+        
+#         # Replace CelebA with your dataset
+#         self.val_dataset = MyCelebA(
+#             self.data_dir,
+#             split='test',
+#             transform=val_transforms,
+#             download=False,
+#         )
 #       ===============================================================
 
-
     def train_dataloader(self) -> DataLoader:
         return DataLoader(
             self.train_dataset,
@@ -185,4 +177,5 @@ def test_dataloader(self) -> Union[DataLoader, List[DataLoader]]:
             num_workers=self.num_workers,
             shuffle=True,
             pin_memory=self.pin_memory,
-        )
+        )
+     

models/mssim_vae.py

@@ -231,7 +231,7 @@ def ssim(self,
         sigma2_sq = F.conv2d(img2 * img2, window, padding = window_size//2, groups=in_channel) - mu2_sq
         sigma12   = F.conv2d(img1 * img2, window, padding = window_size//2, groups=in_channel) - mu1_mu2
 
-        img_range = img1.max() - img1.min()
+        img_range = 1.0 #img1.max() - img1.min() # Dynamic range
         C1 = (0.01 * img_range) ** 2
         C2 = (0.03 * img_range) ** 2
 

requirements.txt

@@ -1,6 +1,6 @@
 pytorch-lightning==1.5.6
 PyYAML==6.0
-tensorboard==2.1.0
+tensorboard>=2.2.0
 torch>=1.6.1
 torchsummary==1.5.1
-torchvision>=0.11.2
+torchvision>=0.10.1

run.py

@@ -29,7 +29,6 @@
         print(exc)
 
 
-
 tb_logger =  TensorBoardLogger(save_dir=config['logging_params']['save_dir'],
                                name=config['model_params']['name'],)
 
@@ -42,6 +41,7 @@
 
 data = VAEDataset(**config["data_params"], pin_memory=len(config['trainer_params']['gpus']) != 0)
 
+data.setup()
 runner = Trainer(logger=tb_logger,
                  callbacks=[
                      LearningRateMonitor(),
@@ -57,5 +57,6 @@
 Path(f"{tb_logger.log_dir}/Samples").mkdir(exist_ok=True, parents=True)
 Path(f"{tb_logger.log_dir}/Reconstructions").mkdir(exist_ok=True, parents=True)
 
+
 print(f"======= Training {config['model_params']['name']} =======")
 runner.fit(experiment, datamodule=data)