Update app.py

app.py

@@ -10,7 +10,57 @@ from typing import Dict
 import functools
 import inspect
 from types import SimpleNamespace
+import torch
+from torch.utils.data import Dataset
+from torchvision import transforms
+import rasterio
+from pathlib import Path
+from torchvision.transforms import ToPILImage
+import numpy as np
+
+class UAHiRISEDataset(Dataset):
+    def __init__(self, root, stage, transform=None):
+        self.root = Path(root)
+        self.stage = stage
+        self.transform = transform
+        self.filenames = self._read_split()
+
+    def __len__(self):
+        return len(self.filenames)
+
+    def __getitem__(self, idx):
+        filename = self.filenames[idx]
+        raster_path = self.root / filename
+
+        raster = rasterio.open(raster_path)
+
+        left = raster.read(1).astype('uint8')
+        dtm = raster.read(2)
 
+        # converting absolute heigths to relative depths
+        dtm = abs(dtm - dtm.min())
+
+        to_pil = ToPILImage()
+
+        to_transform = {"image": to_pil(left).convert('RGB'), "dtm": dtm}
+
+        return self.transform(to_transform)
+        # return to_transform
+
+    def _add_channels(self, image):
+        img_expanded = np.stack([image, image, image], axis=-1)
+        img_tensor = torch.from_numpy(img_expanded).permute(2, 0, 1)
+        return img_tensor
+
+    def set_transform(self, transform):
+        self.transform = transform
+
+    def _read_split(self):
+        split_filename = f'uahirise_{self.stage}.txt'
+        split_filepath = Path(f'filenames/{split_filename}')
+        filenames = split_filepath.read_text().splitlines()
+        return filenames
+        
 class Autoencoder(nn.Module):
   def __init__(self):
     super().__init__()