fix weighting

Author: ccnguyen

interpolate.py

@@ -51,8 +51,8 @@ def _fill(self, holed_img, params):
             idx_select = filled_idx[idx[:, i]]  # size: num_coords, 2
 
             # add value of those coords, weighted by their distance
-            vals += holed_img[idx_select[:, 0], idx_select[:, 1]] * dist[:, i]
-        vals /= torch.sum(dist, dim=1)
+            vals += holed_img[idx_select[:, 0], idx_select[:, 1]] * (1.0 / dist[:, i])
+        vals /= torch.sum((1.0 / dist), dim=1)
         holed_img[unfilled_idx[:, 0], unfilled_idx[:, 1]] = vals
         return holed_img
 
@@ -125,8 +125,8 @@ def _fill(self, holed_img, params):
             idx_select = filled_idx[idx[:, i]]  # size: num_coords, 2
 
             # add value of those coords, weighted by their distance
-            vals += holed_img[idx_select[:, 0], idx_select[:, 1]] * dist[:, i]
-        vals /= torch.sum(dist, dim=1)
+            vals += holed_img[idx_select[:, 0], idx_select[:, 1]] * (1.0 / dist[:, i])
+        vals /= torch.sum((1.0 / dist), dim=1)
         holed_img[unfilled_idx[:, 0], unfilled_idx[:, 1]] = vals
         return holed_img
 

model.py

@@ -1,6 +1,7 @@
 import torch.nn as nn
 from interpolate import TreeMultiQuad, TreeMultiRandom
 import torch
+from debayer import Debayer3x3
 
 
 class TreeModel(nn.Module):
@@ -22,7 +23,7 @@ def __init__(self, coded_type='irregular', sz=512, num_channels=6):
             self.tree = TreeMultiQuad(sz=sz)
         else:
             # irregular spaced points
-            self.tree = TreeMultiRandom(sz=sz, num_channels=6)
+            self.tree = TreeMultiRandom(sz=sz, num_channels=num_channels)
 
     def forward(self, coded, lookup_channels=None):
         ''' Coded is the single channel image we want to stack into multiple channels '''
@@ -32,6 +33,16 @@ def forward(self, coded, lookup_channels=None):
             return self.tree(coded)
 
 
+class InterpModel(nn.Module):
+    ''' Can use a bilinear interpolation module for comparison'''
+    def __init__(self):
+        super().__init__()
+        self.interp = Debayer3x3().cuda()
+
+    def forward(self, coded):
+        return self.interp(coded)
+
+
 if __name__ == '__main__':
     device = 'cuda:0'
     tree = TreeModel(coded_type='irregular', sz=512, num_channels=6)