Major code refactoring

Dataset.py

@@ -1,73 +1,67 @@
-import torch
 import os
+
 import cv2 as cv
 import numpy as np
-import torchvision
-import albumentations as A
-import random
-import einops
-from utilities import augmentations
+import torch
 from tqdm import tqdm
-import matplotlib.pyplot as plt
-from annotation_utils import read_annotation, get_insar_path
+
+from utilities import augmentations
+
 
 class Dataset(torch.utils.data.Dataset):
     def __init__(self, config):
         self.data_path = config["data_path"]
         self.augmentations = augmentations.get_augmentations(config)
         self.interferograms = []
-        self.channels = config['num_channels']
+        self.channels = config["num_channels"]
         frames = os.listdir(self.data_path)
         for frame in tqdm(frames):
-            frame_path = self.data_path + '/'+frame +'/interferograms/'
+            frame_path = self.data_path + "/" + frame + "/interferograms/"
             caption = os.listdir(frame_path)
             for cap in caption:
-                caption_path = frame_path + cap +'/'+cap+ '.geo.diff.png'
+                caption_path = frame_path + cap + "/" + cap + ".geo.diff.png"
 
                 if os.path.exists(caption_path):
-                    image_dict = {'path': caption_path, 'frame': frame}
+                    image_dict = {"path": caption_path, "frame": frame}
                     self.interferograms.append(image_dict)
 
         self.num_examples = len(self.interferograms)
 
-
     def __len__(self):
         return self.num_examples
 
-
-    def prepare_insar(self,insar):
-        insar = torch.from_numpy(insar).float().permute(2,0,1)
+    def prepare_insar(self, insar):
+        insar = torch.from_numpy(insar).float().permute(2, 0, 1)
         insar /= 255
         return insar
 
-    def read_insar(self,path):
-        insar = cv.imread(path,0)
+    def read_insar(self, path):
+        insar = cv.imread(path, 0)
         if insar is None:
-            print('None')
+            print("None")
             return insar
 
-        insar = einops.repeat(insar, 'h w -> h w c', c=self.channels)
+        insar = np.expand_dims(insar, axis=2).repeat(self.channels, axis=2)
         transform = self.augmentations(image=insar)
-        insar_1 = transform['image']
+        insar_1 = transform["image"]
         transform_2 = self.augmentations(image=insar)
-        insar_2 = transform_2['image']
+        insar_2 = transform_2["image"]
 
         insar_1 = self.prepare_insar(insar_1)
         insar_2 = self.prepare_insar(insar_2)
-        return (insar_1,insar_2)
+        return (insar_1, insar_2)
 
     def __getitem__(self, index):
         insar = None
         while insar is None:
             sample = self.interferograms[index]
-            path = sample['path']
+            path = sample["path"]
 
             insar = self.read_insar(path)
             if insar is None:
-                if index<self.num_examples-1:
-                    index +=1
+                if index < self.num_examples - 1:
+                    index += 1
                 else:
                     index = 0
 
-
-        return insar,0
+        return insar, 0

README.md

@@ -1,6 +1,8 @@
 # Hephaestus Dataset
 
-![hephaestus](hephaestus-logo.png)
+<p align="center">
+  <img src="docs/hephaestus-logo.png">
+</p>
 
 This repository contains the data and code used in [Hephaestus: A large scale multitask dataset towards InSAR understanding](https://openaccess.thecvf.com/content/CVPR2022W/EarthVision/papers/Bountos_Hephaestus_A_Large_Scale_Multitask_Dataset_Towards_InSAR_Understanding_CVPRW_2022_paper.pdf) as published in CVPR 2022 workshop Earthvision.
 
@@ -17,9 +19,39 @@ If you use this work, please cite:
 }
 ```
 ### Dependancies
-This repo has been tested with python3.9. To install the necessary dependancies run:
+This repo has been tested with python3.9. To install the necessary dependancies run:   
 `pip install -r requirements.txt`
 
+### Multi-GPU / Multi-Node training
+You can make use of torchrun or SLURM to launch distributed jobs.
+
+###### torchrun   
+Single-Node Multi-GPU:
+```
+torchrun --standalone --nnodes=1 --nproc_per_node=2 main.py
+```
+
+Multi-Node Multi-GPU:
+```
+# On XXX.XXX.XXX.62 (the master node)
+torchrun \
+--nproc_per_node=2 --nnodes=2 --node_rank=0 \
+--master_addr=XXX.XXX.XXX.62 --master_port=1234 \
+main.py
+
+# On XXX.XXX.XXX.63 (the worker node)
+torchrun \
+--nproc_per_node=2 --nnodes=2 --node_rank=1 \
+--master_addr=XXX.XXX.XXX.62 --master_port=1234 \
+main.py
+```
+
+###### SLURM:
+After setting the relevant parameters inside hephaestus.slurm:
+```
+sbatch hephaestus.slurm
+```
+
 ### Dataset and pretrained models
 
 The annotation files can be downloaded [here](https://www.dropbox.com/s/i08mz5514gczksz/annotations_hephaestus.zip?dl=0).
@@ -38,7 +70,7 @@ The dataset is organized in the following structure:
 The cropped 224x224 patches, along with the respective masks and labels can be found [here](https://www.dropbox.com/s/2bkpj79jepk0vks/Hephaestus_Classification.zip?dl=0).
 Some examples of these patches can be seen in the following figure. 
 
-![figure](examples.png)
+![figure](docs/examples.png)
 
 The directory structure for the cropped patches is:
 
@@ -72,7 +104,7 @@ The script will automatically create folders for the checkpoints and store the c
 ### Annotation
 
 The dataset contains both labeled and unlabeled data. The labeled part covers 38 frames summing up to 19,919 annotated InSAR.
-The list of the studied volcanoes, along with the temporal distribution of their samples can be seen below. ![below](volcano_distribution.png)
+The list of the studied volcanoes, along with the temporal distribution of their samples can be seen below. ![below](docs/volcano_distribution.png)
 
 Each labeled InSAR is accompanied by a json file containing the annotation details. Below we present an example of an annotation file. A detailed description can be seen in the original paper (section 2.2):
 ```python