Working demo

.gitattributes

@@ -32,3 +32,70 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+checkpoints/sid_1-500_m3lteacher.pth filter=lfs diff=lfs merge=lfs -text
+checkpoints/sid_1-500_mtteacher.pth filter=lfs diff=lfs merge=lfs -text
+examples/rgb/camera_0b2396756adf4a76bbb985af92f534af_office_7_frame_10_domain.png filter=lfs diff=lfs merge=lfs -text
+examples/rgb/camera_0b2396756adf4a76bbb985af92f534af_office_7_frame_9_domain.png filter=lfs diff=lfs merge=lfs -text
+examples/rgb/camera_0a529ee09fc04037831902b83073862b_office_4_frame_46_domain.png filter=lfs diff=lfs merge=lfs -text
+examples/rgb/camera_0b2396756adf4a76bbb985af92f534af_office_7_frame_4_domain.png filter=lfs diff=lfs merge=lfs -text
+examples/rgb/camera_0c2610e800274e329f5fac53402855ef_office_32_frame_8_domain.png filter=lfs diff=lfs merge=lfs -text
+examples/rgb/camera_0a529ee09fc04037831902b83073862b_office_4_frame_53_domain.png filter=lfs diff=lfs merge=lfs -text
+examples/rgb/camera_0a529ee09fc04037831902b83073862b_office_4_frame_49_domain.png filter=lfs diff=lfs merge=lfs -text
+examples/rgb/camera_0b2396756adf4a76bbb985af92f534af_office_7_frame_8_domain.png filter=lfs diff=lfs merge=lfs -text
+examples/rgb/camera_0c2610e800274e329f5fac53402855ef_office_32_frame_5_domain.png filter=lfs diff=lfs merge=lfs -text
+examples/rgb/camera_4e7d4edf7c2a40ae93437d1fbde22043_office_6_frame_47_domain.png filter=lfs diff=lfs merge=lfs -text
+examples/rgb/camera_0a529ee09fc04037831902b83073862b_office_4_frame_47_domain.png filter=lfs diff=lfs merge=lfs -text
+examples/rgb/camera_0b2396756adf4a76bbb985af92f534af_office_7_frame_1_domain.png filter=lfs diff=lfs merge=lfs -text
+examples/rgb/camera_0c2610e800274e329f5fac53402855ef_office_32_frame_2_domain.png filter=lfs diff=lfs merge=lfs -text
+examples/rgb/camera_0c2610e800274e329f5fac53402855ef_office_32_frame_3_domain.png filter=lfs diff=lfs merge=lfs -text
+examples/rgb/camera_0c2610e800274e329f5fac53402855ef_office_32_frame_9_domain.png filter=lfs diff=lfs merge=lfs -text
+examples/rgb/camera_0a529ee09fc04037831902b83073862b_office_4_frame_48_domain.png filter=lfs diff=lfs merge=lfs -text
+examples/rgb/camera_0b2396756adf4a76bbb985af92f534af_office_7_frame_7_domain.png filter=lfs diff=lfs merge=lfs -text
+examples/rgb/camera_0c2610e800274e329f5fac53402855ef_office_32_frame_4_domain.png filter=lfs diff=lfs merge=lfs -text
+examples/rgb/camera_0c2610e800274e329f5fac53402855ef_office_32_frame_6_domain.png filter=lfs diff=lfs merge=lfs -text
+examples/rgb/camera_0b2396756adf4a76bbb985af92f534af_office_7_frame_5_domain.png filter=lfs diff=lfs merge=lfs -text
+examples/rgb/camera_0b2396756adf4a76bbb985af92f534af_office_7_frame_0_domain.png filter=lfs diff=lfs merge=lfs -text
+examples/rgb/camera_0b2396756adf4a76bbb985af92f534af_office_7_frame_3_domain.png filter=lfs diff=lfs merge=lfs -text
+examples/rgb/camera_0b2396756adf4a76bbb985af92f534af_office_7_frame_6_domain.png filter=lfs diff=lfs merge=lfs -text
+examples/rgb/camera_0c2610e800274e329f5fac53402855ef_office_32_frame_0_domain.png filter=lfs diff=lfs merge=lfs -text
+examples/rgb/camera_0c2610e800274e329f5fac53402855ef_office_32_frame_10_domain.png filter=lfs diff=lfs merge=lfs -text
+examples/rgb/camera_0c2610e800274e329f5fac53402855ef_office_32_frame_7_domain.png filter=lfs diff=lfs merge=lfs -text
+examples/rgb/camera_0a529ee09fc04037831902b83073862b_office_4_frame_52_domain.png filter=lfs diff=lfs merge=lfs -text
+examples/rgb/camera_0b2396756adf4a76bbb985af92f534af_office_7_frame_2_domain.png filter=lfs diff=lfs merge=lfs -text
+examples/rgb/camera_0c2610e800274e329f5fac53402855ef_office_32_frame_1_domain.png filter=lfs diff=lfs merge=lfs -text
+examples/rgb/camera_0a529ee09fc04037831902b83073862b_office_4_frame_50_domain.png filter=lfs diff=lfs merge=lfs -text
+examples/rgb/camera_0a529ee09fc04037831902b83073862b_office_4_frame_51_domain.png filter=lfs diff=lfs merge=lfs -text
+examples/rgb/camera_0a529ee09fc04037831902b83073862b_office_4_frame_45_domain.png filter=lfs diff=lfs merge=lfs -text
+examples/depth/camera_0a529ee09fc04037831902b83073862b_office_4_frame_49_domain.png filter=lfs diff=lfs merge=lfs -text
+examples/depth/camera_0a529ee09fc04037831902b83073862b_office_4_frame_53_domain.png filter=lfs diff=lfs merge=lfs -text
+examples/depth/camera_0b2396756adf4a76bbb985af92f534af_office_7_frame_3_domain.png filter=lfs diff=lfs merge=lfs -text
+examples/depth/camera_0b2396756adf4a76bbb985af92f534af_office_7_frame_6_domain.png filter=lfs diff=lfs merge=lfs -text
+examples/depth/camera_0c2610e800274e329f5fac53402855ef_office_32_frame_0_domain.png filter=lfs diff=lfs merge=lfs -text
+examples/depth/camera_0b2396756adf4a76bbb985af92f534af_office_7_frame_8_domain.png filter=lfs diff=lfs merge=lfs -text
+examples/depth/camera_0c2610e800274e329f5fac53402855ef_office_32_frame_3_domain.png filter=lfs diff=lfs merge=lfs -text
+examples/depth/camera_0c2610e800274e329f5fac53402855ef_office_32_frame_8_domain.png filter=lfs diff=lfs merge=lfs -text
+examples/depth/camera_0a529ee09fc04037831902b83073862b_office_4_frame_46_domain.png filter=lfs diff=lfs merge=lfs -text
+examples/depth/camera_0a529ee09fc04037831902b83073862b_office_4_frame_52_domain.png filter=lfs diff=lfs merge=lfs -text
+examples/depth/camera_0b2396756adf4a76bbb985af92f534af_office_7_frame_5_domain.png filter=lfs diff=lfs merge=lfs -text
+examples/depth/camera_0c2610e800274e329f5fac53402855ef_office_32_frame_4_domain.png filter=lfs diff=lfs merge=lfs -text
+examples/depth/camera_0b2396756adf4a76bbb985af92f534af_office_7_frame_0_domain.png filter=lfs diff=lfs merge=lfs -text
+examples/depth/camera_0b2396756adf4a76bbb985af92f534af_office_7_frame_10_domain.png filter=lfs diff=lfs merge=lfs -text
+examples/depth/camera_4e7d4edf7c2a40ae93437d1fbde22043_office_6_frame_47_domain.png filter=lfs diff=lfs merge=lfs -text
+examples/depth/camera_0a529ee09fc04037831902b83073862b_office_4_frame_47_domain.png filter=lfs diff=lfs merge=lfs -text
+examples/depth/camera_0b2396756adf4a76bbb985af92f534af_office_7_frame_4_domain.png filter=lfs diff=lfs merge=lfs -text
+examples/depth/camera_0c2610e800274e329f5fac53402855ef_office_32_frame_5_domain.png filter=lfs diff=lfs merge=lfs -text
+examples/depth/camera_0c2610e800274e329f5fac53402855ef_office_32_frame_9_domain.png filter=lfs diff=lfs merge=lfs -text
+examples/depth/camera_0a529ee09fc04037831902b83073862b_office_4_frame_48_domain.png filter=lfs diff=lfs merge=lfs -text
+examples/depth/camera_0a529ee09fc04037831902b83073862b_office_4_frame_50_domain.png filter=lfs diff=lfs merge=lfs -text
+examples/depth/camera_0a529ee09fc04037831902b83073862b_office_4_frame_51_domain.png filter=lfs diff=lfs merge=lfs -text
+examples/depth/camera_0b2396756adf4a76bbb985af92f534af_office_7_frame_1_domain.png filter=lfs diff=lfs merge=lfs -text
+examples/depth/camera_0b2396756adf4a76bbb985af92f534af_office_7_frame_9_domain.png filter=lfs diff=lfs merge=lfs -text
+examples/depth/camera_0c2610e800274e329f5fac53402855ef_office_32_frame_10_domain.png filter=lfs diff=lfs merge=lfs -text
+examples/depth/camera_0c2610e800274e329f5fac53402855ef_office_32_frame_6_domain.png filter=lfs diff=lfs merge=lfs -text
+examples/depth/camera_0c2610e800274e329f5fac53402855ef_office_32_frame_2_domain.png filter=lfs diff=lfs merge=lfs -text
+examples/depth/camera_0a529ee09fc04037831902b83073862b_office_4_frame_45_domain.png filter=lfs diff=lfs merge=lfs -text
+examples/depth/camera_0b2396756adf4a76bbb985af92f534af_office_7_frame_2_domain.png filter=lfs diff=lfs merge=lfs -text
+examples/depth/camera_0b2396756adf4a76bbb985af92f534af_office_7_frame_7_domain.png filter=lfs diff=lfs merge=lfs -text
+examples/depth/camera_0c2610e800274e329f5fac53402855ef_office_32_frame_1_domain.png filter=lfs diff=lfs merge=lfs -text
+examples/depth/camera_0c2610e800274e329f5fac53402855ef_office_32_frame_7_domain.png filter=lfs diff=lfs merge=lfs -text
+classcolors.png filter=lfs diff=lfs merge=lfs -text

README.md

@@ -1,7 +1,7 @@
 ---
 title: M3L
-emoji: 🐢
-colorFrom: pink
+emoji: 📚
+colorFrom: purple
 colorTo: gray
 sdk: gradio
 sdk_version: 3.23.0

app.py

@@ -0,0 +1,161 @@
+import torch
+import torch.nn as nn
+import gradio as gr
+import numpy as np
+import os
+import random
+import pickle as pkl
+from models.segmentation_models.linearfusemaskedconsmixbatch.segformer import LinearFusionMaskedConsistencyMixBatch
+from models.segmentation_models.linearfuse.segformer import WeTrLinearFusion
+from datasets.preprocessors import RGBDValPre
+from utils.constants import Constants as C
+
+class Arguments:
+    def __init__(self, ratio):
+        self.ratio = ratio
+        self.masking_ratio = 1.0
+
+colors = pkl.load(open('./colors.pkl', 'rb'))
+args = Arguments(ratio = 0.8)
+
+mtmodel = WeTrLinearFusion("mit_b2", args, num_classes=13, pretrained=False)
+mtmodelpath = './checkpoints/sid_1-500_mtteacher.pth'
+mtmodel.load_state_dict(torch.load(mtmodelpath,  map_location=torch.device('cpu')))
+mtmodel.eval()
+
+m3lmodel = LinearFusionMaskedConsistencyMixBatch("mit_b2", args, num_classes=13, pretrained=False)
+m3lmodelpath = './checkpoints/sid_1-500_m3lteacher.pth'
+m3lmodel.load_state_dict(torch.load(m3lmodelpath,  map_location=torch.device('cpu')))
+m3lmodel.eval()
+
+
+
+class MaskStudentTeacher(nn.Module):
+    
+    def __init__(self, student, teacher, ema_alpha, mode = 'train'):
+        super(MaskStudentTeacher, self).__init__()
+        self.student = student
+        self.teacher = teacher
+        self.teacher = self._detach_teacher(self.teacher)
+        self.ema_alpha = ema_alpha
+        self.mode = mode
+    def forward(self, data, student = True, teacher = True, mask = False, range_batches_to_mask = None, **kwargs):
+        ret = []
+        if student:
+            if self.mode == 'train':
+                ret.append(self.student(data, mask = mask, range_batches_to_mask = range_batches_to_mask, **kwargs))
+            elif self.mode == 'val':
+                ret.append(self.student(data, mask = False, **kwargs))
+            else:
+                raise Exception('Mode not supported')
+        if teacher:
+            ret.append(self.teacher(data, mask = False,  **kwargs)) #Not computing loss for teacher ever but passing the results as if loss was also returned
+        return ret
+    def _detach_teacher(self, model):
+        for param in model.parameters():
+            param.detach_()
+        return model
+    def update_teacher_models(self, global_step):
+        alpha = min(1 - 1 / (global_step + 1), self.ema_alpha)
+        for ema_param, param in zip(self.teacher.parameters(), self.student.parameters()):
+            ema_param.data.mul_(alpha).add_(1 - alpha, param.data)
+        return
+    def copy_student_to_teacher(self):
+        for ema_param, param in zip(self.teacher.parameters(), self.student.parameters()):
+            ema_param.data.mul_(0).add_(param.data)
+        return
+    def get_params(self):
+        student_params = self.student.get_params()
+        teacher_params = self.teacher.get_params()
+        return student_params
+    
+
+def preprocess_data(rgb, depth, dataset_settings):
+    #RGB: np.array, RGB
+    #Depth: np.array, minmax normalized, *255
+    preprocess = RGBDValPre(C.pytorch_mean, C.pytorch_std, dataset_settings)
+    rgb, depth = preprocess(rgb, depth)
+    if rgb is not None:
+        rgb = torch.from_numpy(np.ascontiguousarray(rgb)).float()
+    if depth is not None:
+        depth = torch.from_numpy(np.ascontiguousarray(depth)).float()
+    return rgb, depth
+
+
+def visualize(colors, pred, num_classes, dataset_settings):
+    pred = pred.transpose(1, 2, 0)
+    predvis = np.zeros((dataset_settings['orig_height'], dataset_settings['orig_width'], 3))
+    for i in range(num_classes):
+        color = colors[i]
+        predvis = np.where(pred == i, color, predvis)
+    predvis /= 255.0
+    predvis = predvis[:,:,::-1]
+    return predvis
+
+def predict(rgb, depth, check):
+    dataset_settings = {}
+    dataset_settings['image_height'], dataset_settings['image_width'] = 540, 540
+    dataset_settings['orig_height'], dataset_settings['orig_width'] = 540,540
+    
+    rgb, depth = preprocess_data(rgb, depth, dataset_settings)
+    if rgb is not None:
+        rgb = rgb.unsqueeze(dim = 0)
+    if depth is not None:
+        depth = depth.unsqueeze(dim = 0)
+    ret = [None, None, './classcolors.png']
+    if "Mean Teacher" in check:
+        if rgb is None:
+            rgb = torch.zeros_like(depth)
+        if depth is None:
+            depth = torch.zeros_like(rgb)
+        scores = mtmodel([rgb, depth])[2]
+        scores = torch.nn.functional.interpolate(scores, size = (dataset_settings["orig_height"], dataset_settings["orig_width"]), mode = 'bilinear', align_corners = True)
+        prob = scores.detach()
+        _, pred = torch.max(prob, dim=1)
+        pred = pred.numpy()
+        predvis = visualize(colors, pred, num_classes=13, dataset_settings=dataset_settings)
+        ret[0] = predvis
+    if "M3L" in check:
+        mask = False
+        masking_branch = None
+        if rgb is None:
+            mask = True
+            masking_branch = 0
+        if depth is None:
+            mask = True
+            masking_branch = 1
+        scores = m3lmodel([rgb, depth], mask = mask, masking_branch = masking_branch)[2]
+        scores = torch.nn.functional.interpolate(scores, size = (dataset_settings["orig_height"], dataset_settings["orig_width"]), mode = 'bilinear', align_corners = True)
+        prob = scores.detach()
+        _, pred = torch.max(prob, dim=1)
+        pred = pred.numpy()
+        predvis = visualize(colors, pred, num_classes=13, dataset_settings=dataset_settings)
+        ret[1] = predvis
+    
+    return ret
+
+imgs = os.listdir('./examples/rgb')
+random.shuffle(imgs)
+examples = []
+for img in imgs:
+    examples.append([
+        './examples/rgb/'+img, './examples/depth/'+img, ["M3L", "Mean Teacher"]
+    ])
+
+with gr.Blocks() as demo:
+    with gr.Row():
+        rgbinput = gr.Image(label="RGB Input").style(height=256, width=256)
+        depthinput = gr.Image(label="Depth Input").style(height=256, width=256)
+    with gr.Row():
+        modelcheck = gr.CheckboxGroup(["Mean Teacher", "M3L"], label="Predictions from", info="Predict using model trained with:")
+    with gr.Row():
+        submit_btn = gr.Button("Submit")
+    with gr.Row():
+        mtoutput = gr.Image(label="Mean Teacher Output").style(height=384, width=384)
+        m3loutput = gr.Image(label="M3L Output").style(height=384, width=384)
+        classnameouptut = gr.Image(label="Classes").style(height=384, width=384)
+    with gr.Row():
+        examplesRow = gr.Examples(examples=examples, examples_per_page=10, inputs=[rgbinput, depthinput, modelcheck])
+    submit_btn.click(fn = predict, inputs = [rgbinput, depthinput, modelcheck], outputs = [mtoutput, m3loutput, classnameouptut])
+
+demo.launch()

checkpoints/sid_1-500_m3lteacher.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d5a23d7e2697b44b18e368e01353c328b13055a05a1cb0946ffb95b692d6facd
+size 99192724

checkpoints/sid_1-500_mtteacher.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7eb24d6275e15376c40ec526281d550e4842ffa71aaa7af58fea54cbf56c2eeb
+size 99186911

colors.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3c03050cb753b0802781f0ce92893ac22129c15724dd4ece6f4b9b4a352db591
+size 2342

datasets/base_dataset.py

@@ -0,0 +1,128 @@
+import torch
+import torch.utils.data as data
+import numpy as np
+import cv2
+from PIL import Image
+from utils.img_utils import pad_image_to_shape
+
+class BaseDataset(data.Dataset):
+
+    def __init__(self, dataset_settings, mode, unsupervised):
+        self._mode = mode
+        self.unsupervised = unsupervised
+        self._rgb_path = dataset_settings['rgb_root']
+        self._depth_path = dataset_settings['depth_root']
+        self._gt_path = dataset_settings['gt_root']
+        self._train_source = dataset_settings['train_source']
+        self._eval_source = dataset_settings['eval_source']
+        self.modalities = dataset_settings['modalities']
+        # self._file_length = dataset_settings['max_samples'] 
+        self._required_length = dataset_settings['required_length'] 
+        self._file_names = self._get_file_names(mode)
+        self.model_input_shape = (dataset_settings['image_height'], dataset_settings['image_width'])
+        
+    def __len__(self):
+        if self._required_length is not None:
+            return self._required_length
+        return len(self._file_names) # when model == "val"
+
+    def _get_file_names(self, mode):
+        assert mode in ['train', 'val']
+        source = self._train_source
+        if mode == "val":
+            source = self._eval_source
+
+        file_names = []
+        with open(source) as f:
+            files = f.readlines()
+
+        for item in files:
+            names = self._process_item_names(item)
+            file_names.append(names)
+        
+        if mode == "val":
+            return file_names
+        elif self._required_length <= len(file_names):
+            return file_names[:self._required_length]
+        else:
+            return self._construct_new_file_names(file_names, self._required_length)
+
+    def _construct_new_file_names(self, file_names, length):
+        assert isinstance(length, int)
+        files_len = len(file_names)
+
+        new_file_names = file_names * (length // files_len) #length % files_len items remaining
+
+        rand_indices = torch.randperm(files_len).tolist()
+        new_indices = rand_indices[:length % files_len]
+
+        new_file_names += [file_names[i] for i in new_indices]
+
+        return new_file_names
+
+    def _process_item_names(self, item):
+        item = item.strip()
+        item = item.split('\t')
+        num_modalities = len(self.modalities)
+        num_items = len(item)
+        names = {}
+        if not self.unsupervised:
+            assert num_modalities + 1 == num_items, f"Number of modalities and number of items in file name don't match, len(modalities) = {num_modalities} and len(item) = {num_items}" + item[0]
+            for i, modality in enumerate(self.modalities):
+                names[modality] = item[i]
+            names['gt'] = item[-1]
+        else:
+            assert num_modalities == num_items, f"Number of modalities and number of items in file name don't match, len(modalities) = {num_modalities} and len(item) = {num_items}"
+            for i, modality in enumerate(self.modalities):
+                names[modality] = item[i]
+            names['gt'] = None
+
+        return names
+
+    def _open_rgb(self, rgb_path, dtype = None):
+        bgr = cv2.imread(rgb_path, cv2.IMREAD_COLOR) #cv2 reads in BGR format, HxWxC
+        rgb = np.array(cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB), dtype=dtype) #Pretrained PyTorch model accepts image in RGB
+        return rgb
+
+    def _open_depth(self, depth_path, dtype = None): #returns in HxWx3 with the same image in all channels
+        img_arr = np.array(Image.open(depth_path))
+        if len(img_arr.shape) == 2:  # grayscale
+            img_arr = np.array(np.tile(img_arr, [3, 1, 1]).transpose(1, 2, 0), dtype = dtype)
+        img_arr = (img_arr - img_arr.min()) * 255.0 / (img_arr.max() - img_arr.min())
+        return img_arr
+    
+    def _open_depth_tf_nyu(self, depth_path, dtype = None): #returns in HxWx3 with the same image in all channels
+        img_arr = np.array(Image.open(depth_path))
+        if len(img_arr.shape) == 2:  # grayscale
+            img_arr = np.tile(img_arr, [3, 1, 1]).transpose(1, 2, 0)
+        return img_arr
+
+    def _open_gt(self, gt_path, dtype = None):
+        return np.array(cv2.imread(gt_path,  cv2.IMREAD_GRAYSCALE), dtype=dtype)
+
+    def slide_over_image(self, img, crop_size, stride_rate):
+        H, W, C = img.shape
+        long_size = H if H > W else W
+        output = []
+        if long_size <= min(crop_size[0], crop_size[1]):
+            raise Exception("Crop size is greater than the image size itself. Not handeled right now")
+        
+        else:
+            stride_0 = int(np.ceil(crop_size[0] * stride_rate))
+            stride_1 = int(np.ceil(crop_size[1] * stride_rate))
+            r_grid = int(np.ceil((H - crop_size[0]) / stride_0)) + 1
+            c_grid = int(np.ceil((W - crop_size[1]) / stride_1)) + 1
+
+            for grid_yidx in range(r_grid):
+                for grid_xidx in range(c_grid):
+                    s_x = grid_xidx * stride_1
+                    s_y = grid_yidx * stride_0
+                    e_x = min(s_x + crop_size[1], W)
+                    e_y = min(s_y + crop_size[0], H)
+                    s_x = e_x - crop_size[1]
+                    s_y = e_y - crop_size[0]
+                    img_sub = img[s_y:e_y, s_x: e_x, :]
+                    img_sub, margin = pad_image_to_shape(img_sub, crop_size, cv2.BORDER_CONSTANT, value=0)
+                    output.append((img_sub, np.array([s_y, e_y, s_x, e_x]), margin))
+        
+        return output

datasets/citysunrgbd.py

@@ -0,0 +1,67 @@
+import torch
+import numpy as np
+
+from datasets.base_dataset import BaseDataset
+
+
+class CityScapesSunRGBD(BaseDataset):
+
+    def __init__(self, dataset_settings, mode, unsupervised, preprocess, sliding = False, stride_rate = None):
+        super(CityScapesSunRGBD, self).__init__(dataset_settings, mode, unsupervised)
+        self.preprocess = preprocess
+        self.sliding = sliding
+        self.stride_rate = stride_rate
+        if self.sliding and self._mode == 'train':
+            print("Ensure correct preprocessing is being done!")
+
+    def __getitem__(self, index):
+        # if self._file_length is not None:
+        #     names = self._construct_new_file_names(self._file_length)[index]
+        # else:
+        #     names = self._file_names[index]
+        names = self._file_names[index]
+        rgb_path = self._rgb_path+names['rgb'] 
+        depth_path = self._rgb_path+names['depth'] 
+        if not self.unsupervised:
+            gt_path = self._gt_path+names['gt']
+        item_name = names['rgb'].split("/")[-1].split(".")[0]
+
+        rgb = self._open_rgb(rgb_path)
+        depth = self._open_depth(depth_path)
+        gt = None
+        if not self.unsupervised:
+            gt = self._open_gt(gt_path)
+
+        if not self.sliding:
+            if self.preprocess is not None:
+                rgb, depth, gt = self.preprocess(rgb, depth, gt)
+
+            if self._mode in ['train', 'val']:
+                rgb = torch.from_numpy(np.ascontiguousarray(rgb)).float()
+                depth = torch.from_numpy(np.ascontiguousarray(depth)).float()
+                if gt is not None:
+                    gt = torch.from_numpy(np.ascontiguousarray(gt)).long()
+            else:
+                raise Exception(f"{self._mode} not supported in CityScapesSunRGB")
+                
+            # output_dict = dict(rgb=rgb, fn=str(item_name),
+            #                    n=len(self._file_names))
+            output_dict = dict(data=[rgb, depth], name = item_name)
+            if gt is not None:
+                output_dict['gt'] = gt
+            return output_dict
+    
+        else:
+            sliding_ouptut = self.slide_over_image(rgb, self.model_input_shape, self.stride_rate)
+            output_dict = {}
+            if self._mode in ['train', 'val']:
+                if gt is not None:
+                    gt = torch.from_numpy(np.ascontiguousarray(gt)).long()
+                output_dict['gt'] = gt
+                output_dict['sliding_output'] = []
+                for img_sub, pos, margin in sliding_ouptut:
+                    if self.preprocess is not None:
+                        img_sub, _ = self.preprocess(img_sub, None)
+                    img_sub = torch.from_numpy(np.ascontiguousarray(img_sub)).float()
+                    output_dict['sliding_output'].append(([img_sub], pos, margin))
+            return output_dict

datasets/get_dataset.py

@@ -0,0 +1,146 @@
+import torch
+from torch.utils.data import DataLoader
+from datasets.citysundepth import CityScapesSunDepth
+from datasets.citysunrgb import CityScapesSunRGB
+from datasets.citysunrgbd import CityScapesSunRGBD
+from datasets.preprocessors import DepthTrainPre, DepthValPre, NYURGBDTrainPre, NYURGBDValPre, RGBDTrainPre, RGBDValPre, RGBTrainPre, RGBValPre
+from datasets.tfnyu import TFNYU
+from utils.constants import Constants as C
+
+def get_dataset(args):
+    datasetClass = None
+    if args.data == "nyudv2":
+        return TFNYU
+    if args.data == "city" or args.data == "sunrgbd" or args.data == 'stanford_indoor':
+        if len(args.modalities) == 1 and args.modalities[0] == 'rgb':
+            datasetClass = CityScapesSunRGB
+        elif len(args.modalities) == 1 and args.modalities[0] == 'depth':
+            datasetClass = CityScapesSunDepth
+        elif len(args.modalities) == 2 and args.modalities[0] == 'rgb' and args.modalities[1] == 'depth':
+            datasetClass = CityScapesSunRGBD
+        else:
+            raise Exception(f"{args.modalities} not configured in get_dataset function.")
+    else:
+        raise Exception(f"{args.data} not configured in get_dataset function.")
+    return datasetClass
+
+def get_preprocessors(args, dataset_settings, mode):
+    if args.data == "nyudv2" and len(args.modalities) == 2 and args.modalities[0] == 'rgb' and args.modalities[1] == 'depth':
+        if mode == 'train':
+            return NYURGBDTrainPre(C.pytorch_mean, C.pytorch_std, dataset_settings)
+        elif mode == 'val':
+            return NYURGBDValPre(C.pytorch_mean, C.pytorch_std, dataset_settings)
+
+    if len(args.modalities) == 1 and args.modalities[0] == 'rgb':
+        if mode == 'train':
+            return RGBTrainPre(C.pytorch_mean, C.pytorch_std, dataset_settings)
+        elif mode == 'val':
+            return RGBValPre(C.pytorch_mean, C.pytorch_std, dataset_settings)
+        else:
+            return Exception("%s mode not defined" % mode)
+    elif len(args.modalities) == 1 and args.modalities[0] == 'depth':
+        if mode == 'train':
+            return DepthTrainPre(dataset_settings)
+        elif mode == 'val':
+            return DepthValPre(dataset_settings)
+        else:
+            return Exception("%s mode not defined" % mode)
+    elif len(args.modalities) == 2 and args.modalities[0] == 'rgb' and args.modalities[1] == 'depth':
+        if mode == 'train':
+            return RGBDTrainPre(C.pytorch_mean, C.pytorch_std, dataset_settings)
+        elif mode == 'val':
+            return RGBDValPre(C.pytorch_mean, C.pytorch_std, dataset_settings)
+        else:
+            return Exception("%s mode not defined" % mode)
+    else:
+        raise Exception("%s not configured for preprocessing" % args.modalities)
+
+def get_train_loader(datasetClass, args, train_source, unsupervised = False):
+    dataset_settings = {'rgb_root': args.rgb_root,
+                        'gt_root': args.gt_root,
+                        'depth_root': args.depth_root,
+                        'train_source': train_source,
+                        'eval_source': args.eval_source, 
+                        'required_length': args.total_train_imgs, #Every dataloader will have  Total Train Images / batch size iterations to be consistent
+                        # 'max_samples': args.max_samples, #Every dataloader will have  Total Train Images / batch size iterations to be consistent
+                        'train_scale_array': args.train_scale_array,
+                        'image_height': args.image_height,
+                        'image_width': args.image_width,
+                        'modalities': args.modalities}
+
+    preprocessing = get_preprocessors(args, dataset_settings, "train")
+    train_dataset = datasetClass(dataset_settings, "train", unsupervised, preprocessing)
+    train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset, num_replicas = args.world_size, rank = args.rank)
+    if unsupervised and "unsup_batch_size" in args:
+        batch_size = args.unsup_batch_size
+    else:
+        batch_size = args.batch_size
+    train_loader = DataLoader(train_dataset, 
+                                batch_size = args.batch_size // args.world_size,
+                                num_workers = args.num_workers, 
+                                drop_last = True, 
+                                shuffle = False,
+                                sampler = train_sampler)
+    return train_loader
+    
+def get_val_loader(datasetClass, args):
+    dataset_settings = {'rgb_root': args.rgb_root,
+                        'gt_root': args.gt_root,
+                        'depth_root': args.depth_root,
+                        'train_source': None,
+                        'eval_source': args.eval_source, 
+                        'required_length': None,
+                        'max_samples': None,
+                        'train_scale_array': args.train_scale_array,
+                        'image_height': args.image_height,
+                        'image_width': args.image_width,
+                        'modalities': args.modalities}
+    if args.data == 'sunrgbd':
+        eval_sources = []
+        for shape in ['427_561', '441_591', '530_730', '531_681']:
+            eval_sources.append(dataset_settings['eval_source'].split('.')[0] + '_' +  shape + '.txt')
+    else:
+        eval_sources = [args.eval_source]
+
+    preprocessing = get_preprocessors(args, dataset_settings, "val")
+    if args.sliding_eval:
+        collate_fn = _sliding_collate_fn
+    else:
+        collate_fn = None
+
+    val_loaders = []
+    for eval_source in eval_sources:
+        dataset_settings['eval_source'] = eval_source
+        val_dataset = datasetClass(dataset_settings, "val", False, preprocessing, args.sliding_eval, args.stride_rate)
+        if args.rank is not None: #DDP Evaluation
+            val_sampler = torch.utils.data.distributed.DistributedSampler(val_dataset, num_replicas = args.world_size, rank = args.rank)
+            batch_size = args.val_batch_size // args.world_size
+        else: #DP Evaluation
+            val_sampler = None
+            batch_size = args.val_batch_size
+
+        val_loader = DataLoader(val_dataset,
+                            batch_size = batch_size,
+                            num_workers = 4, 
+                            drop_last = False, 
+                            shuffle = False, 
+                            collate_fn = collate_fn, 
+                            sampler = val_sampler)
+        val_loaders.append(val_loader)
+    return val_loaders
+    
+
+def _sliding_collate_fn(batch):
+        gt = torch.stack([b['gt'] for b in batch])
+        sliding_output = []
+        num_modalities = len(batch[0]['sliding_output'][0][0])
+        for i in range(len(batch[0]['sliding_output'])): #i iterates over positions
+            imgs = [torch.stack([b['sliding_output'][i][0][m] for b in batch]) for m in range(num_modalities)]
+            pos = batch[0]['sliding_output'][i][1]
+            pos_compare = [(b['sliding_output'][i][1] == pos).all() for b in batch]
+            assert all(pos_compare), f"Position not same for all points in the batch: {pos_compare}, {[b['sliding_output'][i][1] for b in batch]}"
+            margin = batch[0]['sliding_output'][i][2]
+            margin_compare = [(b['sliding_output'][i][2] == margin).all() for b in batch]
+            assert all(margin_compare), f"Margin not same for all points in the batch: {margin_compare}, {[b['sliding_output'][i][2] for b in batch]}"
+            sliding_output.append((imgs, pos, margin))
+        return {"gt": gt, "sliding_output": sliding_output}

datasets/preprocessors.py

@@ -0,0 +1,144 @@
+from utils.img_utils import normalizedepth, random_crop_pad_to_shape, random_mirror, random_scale, normalize, resizedepth, resizergb, tfnyu_normalizedepth
+
+class RGBTrainPre(object):
+    def __init__(self, pytorch_mean, pytorch_std, dataset_settings):
+        self.pytorch_mean = pytorch_mean
+        self.pytorch_std = pytorch_std
+        self.train_scale_array = dataset_settings['train_scale_array']
+        self.crop_size = (dataset_settings['image_height'], dataset_settings['image_width'])
+
+    def __call__(self, rgb, gt):
+        transformed_dict = random_mirror({"rgb":rgb, "gt":gt})
+        if self.train_scale_array is not None:
+            transformed_dict, _ = random_scale(transformed_dict, self.train_scale_array, (rgb.shape[0], rgb.shape[1]))
+
+        transformed_dict, _ = random_crop_pad_to_shape(transformed_dict, transformed_dict['rgb'].shape[:2], self.crop_size) #Makes gt HxWx1
+        rgb = transformed_dict['rgb']
+        gt = transformed_dict['gt']
+        rgb = normalize(rgb, self.pytorch_mean, self.pytorch_std)
+
+        rgb = rgb.transpose(2, 0, 1) #Brings the channel dimension in the top. Final output = CxHxW
+        return rgb, gt 
+
+
+class RGBValPre(object):
+    def __init__(self, pytorch_mean, pytorch_std, dataset_settings):
+        self.pytorch_mean = pytorch_mean
+        self.pytorch_std = pytorch_std
+        self.model_input_shape = (dataset_settings['image_height'], dataset_settings['image_width'])
+
+    def __call__(self, rgb, gt):
+        rgb = resizergb(rgb, self.model_input_shape)
+        rgb = normalize(rgb, self.pytorch_mean, self.pytorch_std)
+        rgb = rgb.transpose(2, 0, 1) #Brings the channel dimension in the top. Final output = CxHxW
+        return rgb, gt
+
+
+class RGBDTrainPre(object):
+    def __init__(self, pytorch_mean, pytorch_std, dataset_settings):
+        self.pytorch_mean = pytorch_mean
+        self.pytorch_std = pytorch_std
+        self.train_scale_array = dataset_settings['train_scale_array']
+        self.crop_size = (dataset_settings['image_height'], dataset_settings['image_width'])
+
+    def __call__(self, rgb, depth, gt):
+        transformed_dict = random_mirror({"rgb":rgb, "depth": depth, "gt":gt})
+        if self.train_scale_array is not None:
+            transformed_dict, _ = random_scale(transformed_dict, self.train_scale_array, (rgb.shape[0], rgb.shape[1]))
+
+        transformed_dict, _ = random_crop_pad_to_shape(transformed_dict, transformed_dict['rgb'].shape[:2], self.crop_size) #Makes gt HxWx1
+        rgb = transformed_dict['rgb']
+        depth = transformed_dict['depth']
+        gt = transformed_dict['gt']
+        rgb = normalize(rgb, self.pytorch_mean, self.pytorch_std)
+        depth = normalizedepth(depth)
+        rgb = rgb.transpose(2, 0, 1) #Brings the channel dimension in the top. Final output = CxHxW
+        depth = depth.transpose(2, 0, 1) #Brings the channel dimension in the top. Final output = CxHxW
+        return rgb, depth, gt 
+
+
+class RGBDValPre(object):
+    def __init__(self, pytorch_mean, pytorch_std, dataset_settings):
+        self.pytorch_mean = pytorch_mean
+        self.pytorch_std = pytorch_std
+        self.model_input_shape = (dataset_settings['image_height'], dataset_settings['image_width'])
+
+    def __call__(self, rgb, depth):
+        if rgb is not None:
+            rgb = resizergb(rgb, self.model_input_shape)
+            rgb = normalize(rgb, self.pytorch_mean, self.pytorch_std)
+            rgb = rgb.transpose(2, 0, 1) #Brings the channel dimension in the top. Final output = CxHxW
+        if depth is not None:
+            depth = resizedepth(depth, self.model_input_shape)
+            depth = normalizedepth(depth)
+            depth = depth.transpose(2, 0, 1) #Brings the channel dimension in the top. Final output = CxHxW
+        
+        return rgb, depth
+
+
+class NYURGBDTrainPre(object):
+    def __init__(self, pytorch_mean, pytorch_std, dataset_settings):
+        self.pytorch_mean = pytorch_mean
+        self.pytorch_std = pytorch_std
+        self.train_scale_array = dataset_settings['train_scale_array']
+        self.crop_size = (dataset_settings['image_height'], dataset_settings['image_width'])
+
+    def __call__(self, rgb, depth, gt):
+        transformed_dict = random_mirror({"rgb":rgb, "depth": depth, "gt":gt})
+        if self.train_scale_array is not None:
+            transformed_dict, _ = random_scale(transformed_dict, self.train_scale_array, (rgb.shape[0], rgb.shape[1]))
+
+        transformed_dict, _ = random_crop_pad_to_shape(transformed_dict, transformed_dict['rgb'].shape[:2], self.crop_size) #Makes gt HxWx1
+        rgb = transformed_dict['rgb']
+        depth = transformed_dict['depth']
+        gt = transformed_dict['gt']
+        rgb = normalize(rgb, self.pytorch_mean, self.pytorch_std)
+        depth = tfnyu_normalizedepth(depth)
+        rgb = rgb.transpose(2, 0, 1) #Brings the channel dimension in the top. Final output = CxHxW
+        depth = depth.transpose(2, 0, 1) #Brings the channel dimension in the top. Final output = CxHxW
+        return rgb, depth, gt 
+
+
+class NYURGBDValPre(object):
+    def __init__(self, pytorch_mean, pytorch_std, dataset_settings):
+        self.pytorch_mean = pytorch_mean
+        self.pytorch_std = pytorch_std
+        self.model_input_shape = (dataset_settings['image_height'], dataset_settings['image_width'])
+
+    def __call__(self, rgb, depth, gt):
+        rgb = resizergb(rgb, self.model_input_shape)
+        depth = resizedepth(depth, self.model_input_shape)
+        rgb = normalize(rgb, self.pytorch_mean, self.pytorch_std)
+        depth = tfnyu_normalizedepth(depth)
+        rgb = rgb.transpose(2, 0, 1) #Brings the channel dimension in the top. Final output = CxHxW
+        depth = depth.transpose(2, 0, 1) #Brings the channel dimension in the top. Final output = CxHxW
+        return rgb, depth, gt
+
+
+class DepthTrainPre(object):
+    def __init__(self, dataset_settings):
+        self.train_scale_array = dataset_settings['train_scale_array']
+        self.crop_size = (dataset_settings['image_height'], dataset_settings['image_width'])
+
+    def __call__(self, depth, gt):
+        transformed_dict = random_mirror({"depth": depth, "gt":gt})
+        if self.train_scale_array is not None:
+            transformed_dict, _ = random_scale(transformed_dict, self.train_scale_array, (depth.shape[0], depth.shape[1]))
+
+        transformed_dict, _ = random_crop_pad_to_shape(transformed_dict, transformed_dict['depth'].shape[:2], self.crop_size) #Makes gt HxWx1
+        depth = transformed_dict['depth']
+        gt = transformed_dict['gt']
+        depth = normalizedepth(depth)
+        depth = depth.transpose(2, 0, 1) #Brings the channel dimension in the top. Final output = CxHxW
+        return depth, gt 
+
+
+class DepthValPre(object):
+    def __init__(self, dataset_settings):
+        self.model_input_shape = (dataset_settings['image_height'], dataset_settings['image_width'])
+
+    def __call__(self, depth, gt):
+        depth = resizedepth(depth, self.model_input_shape)
+        depth = normalizedepth(depth)
+        depth = depth.transpose(2, 0, 1) #Brings the channel dimension in the top. Final output = CxHxW
+        return depth, gt