MCPcopy Create free account
hub / github.com/drinkingcoder/NeuralMarker / convert_flow_to_mapping

Function convert_flow_to_mapping

core/utils/utils.py:244–314  ·  view source on GitHub ↗
(flow, output_channel_first=True)

Source from the content-addressed store, hash-verified

242 return grid
243
244def convert_flow_to_mapping(flow, output_channel_first=True):
245 if not isinstance(flow, np.ndarray):
246 # torch tensor
247 if len(flow.shape) == 4:
248 if flow.shape[1] != 2:
249 # size is BxHxWx2
250 flow = flow.permute(0, 3, 1, 2)
251
252 B, C, H, W = flow.size()
253
254 xx = torch.arange(0, W).view(1, -1).repeat(H, 1)
255 yy = torch.arange(0, H).view(-1, 1).repeat(1, W)
256 xx = xx.view(1, 1, H, W).repeat(B, 1, 1, 1)
257 yy = yy.view(1, 1, H, W).repeat(B, 1, 1, 1)
258 grid = torch.cat((xx, yy), 1).float()
259
260 if flow.is_cuda:
261 grid = grid.cuda()
262 map = flow + grid # here also channel first
263 if not output_channel_first:
264 map = map.permute(0,2,3,1)
265 else:
266 if flow.shape[0] != 2:
267 # size is HxWx2
268 flow = flow.permute(2, 0, 1)
269
270 C, H, W = flow.size()
271
272 xx = torch.arange(0, W).view(1, -1).repeat(H, 1)
273 yy = torch.arange(0, H).view(-1, 1).repeat(1, W)
274 xx = xx.view(1, H, W)
275 yy = yy.view(1, H, W)
276 grid = torch.cat((xx, yy), 0).float() # attention, concat axis=0 here
277
278 if flow.is_cuda:
279 grid = grid.cuda()
280 map = flow + grid # here also channel first
281 if not output_channel_first:
282 map = map.permute(1,2,0).float()
283 return map.float()
284 else:
285 # here numpy arrays
286 if len(flow.shape) == 4:
287 if flow.shape[3] != 2:
288 # size is Bx2xHxW
289 flow = flow.transpose(0, 2, 3, 1)
290 # BxHxWx2
291 b, h_scale, w_scale = flow.shape[:3]
292 map = np.copy(flow)
293 X, Y = np.meshgrid(np.linspace(0, w_scale - 1, w_scale),
294 np.linspace(0, h_scale - 1, h_scale))
295 for i in range(b):
296 map[i, :, :, 0] = flow[i, :, :, 0] + X
297 map[i, :, :, 1] = flow[i, :, :, 1] + Y
298 if output_channel_first:
299 map = map.transpose(0,3,1,2)
300 else:
301 if flow.shape[0] == 2:

Callers 1

Calls

no outgoing calls

Tested by

no test coverage detected