lakshjaisinghani / estimating-depth-from-rgb-and-sparse-sensing Goto Github PK

model = D3().float()
model = model.to(device)
Loss_fun  = nn.MSELoss()
optimizer = optim.Adam(net.parameters(), lr = learning_rate)

optimizer = optim.Adam(model.parameters(), lr = learning_rate)

net = D3().float()
net = net.to(device)

class NYU_V2(Dataset):
  def __init__(self, trn_tst=0, transform=None):
    data = h5py.File('./nyu_depth_v2_labeled.mat')
    
    if trn_tst == 0:
      #trainloader
      self.images = data["images"][0:1400]
      self.depths  = data["depths"][0:1400]
    else:
      #testloader
      self.images = data["images"][1401:]
      self.depths = data["depths"][1401:]
    
    self.transform = transform
    
  def __len__(self):
    return len(self.images)
 
  def __getitem__(self, idx):
    if torch.is_tensor(idx):
      idx = idx.tolist()
      
    sample = self.images[idx, :]
    s_depth = self.depths[idx, :]
    sample = torch.from_numpy(np.transpose(sample, (2, 1, 0)))
    s_depth=  torch.from_numpy(np.transpose(s_depth, (1, 0)))
    
    return sample.float(), s_depth.float()

class NYU_V2(Dataset):
    def __init__(self, mode="train", transform=None, path="./nyu_depth_v2_labeled.mat"):
        f = h5py.File(path)
        
        self.images = f["images"]
        self.raw_depths = f["rawDepths"]
        self.gt_depths = f["depths"]
        self.transform = transform
        
    def __len__(self):
        return len(self.images)

    def __getitem__(self, idx):
        if torch.is_tensor(idx):
            idx = idx.tolist()
        
        # Data
        img = self.images[idx, :]
        img = torch.from_numpy(np.transpose(img, (2, 1, 0)))
        raw_depth = self.raw_depths[idx, :]
        raw_depth = torch.from_numpy(np.transpose(raw_depth, (1,0)))
        
        # Label
        depth = self.gt_depths[idx, :]
        depth = torch.from_numpy(np.transpose(depth, (1,0)))
        
        return (img.float(), raw_depth.float()), depth.float()