def test_flow_to_image(batch):
h, w = 100, 100
flow = torch.meshgrid(torch.arange(h), torch.arange(w), indexing="ij")
flow = torch.stack(flow[::-1], dim=0).float()
flow[0] -= h / 2
flow[1] -= w / 2
if batch:
flow = torch.stack([flow, flow])
img = utils.flow_to_image(flow)
assert img.shape == (2, 3, h, w) if batch else (3, h, w)
path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "assets",
"expected_flow.pt")
expected_img = torch.load(path, map_location="cpu")
if batch:
expected_img = torch.stack([expected_img, expected_img])
assert_equal(expected_img, img)
def test_flow_to_image_errors(input_flow, match):
with pytest.raises(ValueError, match=match):
utils.flow_to_image(flow=input_flow)
#################################### # Visualizing predicted flows # --------------------------- # Torchvision provides the :func:`~torchvision.utils.flow_to_image` utlity to # convert a flow into an RGB image. It also supports batches of flows. # each "direction" in the flow will be mapped to a given RGB color. In the # images below, pixels with similar colors are assumed by the model to be moving # in similar directions. The model is properly able to predict the movement of # the ball and the player. Note in particular the different predicted direction # of the ball in the first image (going to the left) and in the second image # (going up). from torchvision.utils import flow_to_image flow_imgs = flow_to_image(predicted_flows) # The images have been mapped into [-1, 1] but for plotting we want them in [0, 1] img1_batch = [(img1 + 1) / 2 for img1 in img1_batch] grid = [[img1, flow_img] for (img1, flow_img) in zip(img1_batch, flow_imgs)] plot(grid) #################################### # Bonus: Creating GIFs of predicted flows # --------------------------------------- # In the example above we have only shown the predicted flows of 2 pairs of # frames. A fun way to apply the Optical Flow models is to run the model on an # entire video, and create a new video from all the predicted flows. Below is a # snippet that can get you started with this. We comment out the code, because # this example is being rendered on a machine without a GPU, and it would take