manually merge pull request #3

Author: oeway

torch_deform_conv/deform_conv.py

@@ -66,7 +66,13 @@ def th_map_coordinates(input, coords, order=1):
 
 def sp_batch_map_coordinates(inputs, coords):
     """Reference implementation for batch_map_coordinates"""
-    coords = coords.clip(0, inputs.shape[1] - 1)
+    # coords = coords.clip(0, inputs.shape[1] - 1)
+
+    assert (coords.shape[2] == 2)
+    height = coords[:,:,0].clip(0, inputs.shape[1] - 1)
+    width = coords[:,:,1].clip(0, inputs.shape[2] - 1)
+    np.concatenate((np.expand_dims(height, axis=2), np.expand_dims(width, axis=2)), 2)
+
     mapped_vals = np.array([
         sp_map_coordinates(input, coord.T, mode='nearest', order=1)
         for input, coord in zip(inputs, coords)
@@ -87,15 +93,22 @@ def th_batch_map_coordinates(input, coords, order=1):
     """
 
     batch_size = input.size(0)
-    input_size = input.size(1)
+    input_height = input.size(1)
+    input_width = input.size(2)
+
     n_coords = coords.size(1)
 
-    coords = torch.clamp(coords, 0, input_size - 1)
+    # coords = torch.clamp(coords, 0, input_size - 1)
+
+    coords = torch.cat((torch.clamp(coords.narrow(2, 0, 1), 0, input_height - 1), torch.clamp(coords.narrow(2, 1, 1), 0, input_width - 1)), 2)
+
+    assert (coords.size(1) == n_coords)
+
     coords_lt = coords.floor().long()
     coords_rb = coords.ceil().long()
     coords_lb = torch.stack([coords_lt[..., 0], coords_rb[..., 1]], 2)
     coords_rt = torch.stack([coords_rb[..., 0], coords_lt[..., 1]], 2)
-    idx = th_repeat(torch.range(0, batch_size-1), n_coords).long()
+    idx = th_repeat(torch.arange(0, batch_size), n_coords).long()
     idx = Variable(idx, requires_grad=False)
     if input.is_cuda:
         idx = idx.cuda()
@@ -108,7 +121,7 @@ def _get_vals_by_coords(input, coords):
         vals = th_flatten(input).index_select(0, inds)
         vals = vals.view(batch_size, n_coords)
         return vals
-    
+
     vals_lt = _get_vals_by_coords(input, coords_lt.detach())
     vals_rb = _get_vals_by_coords(input, coords_rb.detach())
     vals_lb = _get_vals_by_coords(input, coords_lb.detach())
@@ -125,21 +138,22 @@ def sp_batch_map_offsets(input, offsets):
     """Reference implementation for tf_batch_map_offsets"""
 
     batch_size = input.shape[0]
-    input_size = input.shape[1]
+    input_height = input.shape[1]
+    input_width = input.shape[2]
 
     offsets = offsets.reshape(batch_size, -1, 2)
-    grid = np.stack(np.mgrid[:input_size, :input_size], -1).reshape(-1, 2)
+    grid = np.stack(np.mgrid[:input_height, :input_width], -1).reshape(-1, 2)
     grid = np.repeat([grid], batch_size, axis=0)
     coords = offsets + grid
-    coords = coords.clip(0, input_size - 1)
+    # coords = coords.clip(0, input_size - 1)
 
     mapped_vals = sp_batch_map_coordinates(input, coords)
     return mapped_vals
 
 
-def th_generate_grid(batch_size, input_size, dtype, cuda):
+def th_generate_grid(batch_size, input_height, input_width, dtype, cuda):
     grid = np.meshgrid(
-        range(input_size), range(input_size), indexing='ij'
+        range(input_height), range(input_width), indexing='ij'
     )
     grid = np.stack(grid, axis=-1)
     grid = grid.reshape(-1, 2)
@@ -162,11 +176,12 @@ def th_batch_map_offsets(input, offsets, grid=None, order=1):
     torch.Tensor. shape = (b, s, s)
     """
     batch_size = input.size(0)
-    input_size = input.size(1)
+    input_height = input.size(1)
+    input_width = input.size(2)
 
     offsets = offsets.view(batch_size, -1, 2)
     if grid is None:
-        grid = th_generate_grid(batch_size, input_size, offsets.data.type(), offsets.data.is_cuda)
+        grid = th_generate_grid(batch_size, input_height, input_width, offsets.data.type(), offsets.data.is_cuda)
 
     coords = offsets + grid
 

torch_deform_conv/layers.py

@@ -54,14 +54,14 @@ def forward(self, x):
 
     @staticmethod
     def _get_grid(self, x):
-        batch_size, input_size= x.size(0), x.size(1)
+        batch_size, input_height, input_width = x.size(0), x.size(1), x.size(2)
         dtype, cuda = x.data.type(), x.data.is_cuda
-        if self._grid_param == (batch_size, input_size, dtype, cuda):
+        if self._grid_param == (batch_size, input_height, input_width, dtype, cuda):
             return self._grid
-        self._grid_param = (batch_size, input_size, dtype, cuda)
-        self._grid = th_generate_grid(batch_size, input_size, dtype, cuda)
+        self._grid_param = (batch_size, input_height, input_width, dtype, cuda)
+        self._grid = th_generate_grid(batch_size, input_height, input_width, dtype, cuda)
         return self._grid
-    
+
     @staticmethod
     def _init_weights(weights, std):
         fan_out = weights.size(0)