| import gradio as gr | |
| """ | |
| ===================================================== | |
| Optical Flow: Predicting movement with the RAFT model | |
| ===================================================== | |
| Optical flow is the task of predicting movement between two images, usually two | |
| consecutive frames of a video. Optical flow models take two images as input, and | |
| predict a flow: the flow indicates the displacement of every single pixel in the | |
| first image, and maps it to its corresponding pixel in the second image. Flows | |
| are (2, H, W)-dimensional tensors, where the first axis corresponds to the | |
| predicted horizontal and vertical displacements. | |
| The following example illustrates how torchvision can be used to predict flows | |
| using our implementation of the RAFT model. We will also see how to convert the | |
| predicted flows to RGB images for visualization. | |
| """ | |
| import numpy as np | |
| import torch | |
| import matplotlib.pyplot as plt | |
| import torchvision.transforms.functional as F | |
| from torchvision.io import read_video | |
| from torchvision.models.optical_flow import Raft_Large_Weights | |
| from torchvision.models.optical_flow import raft_large | |
| import tempfile | |
| from pathlib import Path | |
| from urllib.request import urlretrieve | |
| def infer(): | |
| video_url = "https://download.pytorch.org/tutorial/pexelscom_pavel_danilyuk_basketball_hd.mp4" | |
| video_path = Path(tempfile.mkdtemp()) / "basketball.mp4" | |
| _ = urlretrieve(video_url, video_path) | |
| frames, _, _ = read_video(str(video_path), output_format="TCHW") | |
| img1= frames[100] | |
| img2 = frames[101] | |
| weights = Raft_Large_Weights.DEFAULT | |
| transforms = weights.transforms() | |
| def preprocess(img1, img2): | |
| img1 = F.resize(img1, size=[520, 960]) | |
| img2 = F.resize(img2, size=[520, 960]) | |
| return transforms(img1, img2) | |
| img1, img2 = preprocess(img1, img2) | |
| print(f"shape = {img1.shape}, dtype = {img1.dtype}") | |
| #################################### | |
| # Estimating Optical flow using RAFT | |
| # ---------------------------------- | |
| # We will use our RAFT implementation from | |
| # :func:`~torchvision.models.optical_flow.raft_large`, which follows the same | |
| # architecture as the one described in the `original paper <https://arxiv.org/abs/2003.12039>`_. | |
| # We also provide the :func:`~torchvision.models.optical_flow.raft_small` model | |
| # builder, which is smaller and faster to run, sacrificing a bit of accuracy. | |
| # If you can, run this example on a GPU, it will be a lot faster. | |
| device = "cuda" if torch.cuda.is_available() else "cpu" | |
| model = raft_large(weights=Raft_Large_Weights.DEFAULT, progress=False).to(device) | |
| model = model.eval() | |
| list_of_flows = model(img1.to(device), img2.to(device)) | |
| print(f"type = {type(list_of_flows)}") | |
| print(f"length = {len(list_of_flows)} = number of iterations of the model") | |
| #################################### | |
| # The RAFT model outputs lists of predicted flows where each entry is a | |
| # (N, 2, H, W) batch of predicted flows that corresponds to a given "iteration" | |
| # in the model. For more details on the iterative nature of the model, please | |
| # refer to the `original paper <https://arxiv.org/abs/2003.12039>`_. Here, we | |
| # are only interested in the final predicted flows (they are the most acccurate | |
| # ones), so we will just retrieve the last item in the list. | |
| # | |
| # As described above, a flow is a tensor with dimensions (2, H, W) (or (N, 2, H, | |
| # W) for batches of flows) where each entry corresponds to the horizontal and | |
| # vertical displacement of each pixel from the first image to the second image. | |
| # Note that the predicted flows are in "pixel" unit, they are not normalized | |
| # w.r.t. the dimensions of the images. | |
| predicted_flows = list_of_flows[-1] | |
| print(f"dtype = {predicted_flows.dtype}") | |
| print(f"shape = {predicted_flows.shape} = (N, 2, H, W)") | |
| print(f"min = {predicted_flows.min()}, max = {predicted_flows.max()}") | |
| #################################### | |
| # 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) | |
| print(flow_imgs) | |
| return "done" | |
| #################################### | |
| # 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 | |
| # too long to run it. | |
| # from torchvision.io import write_jpeg | |
| # for i, (img1, img2) in enumerate(zip(frames, frames[1:])): | |
| # # Note: it would be faster to predict batches of flows instead of individual flows | |
| # img1, img2 = preprocess(img1, img2) | |
| # list_of_flows = model(img1.to(device), img2.to(device)) | |
| # predicted_flow = list_of_flows[-1][0] | |
| # flow_img = flow_to_image(predicted_flow).to("cpu") | |
| # output_folder = "/tmp/" # Update this to the folder of your choice | |
| # write_jpeg(flow_img, output_folder + f"predicted_flow_{i}.jpg") | |
| #################################### | |
| # Once the .jpg flow images are saved, you can convert them into a video or a | |
| # GIF using ffmpeg with e.g.: | |
| # | |
| # ffmpeg -f image2 -framerate 30 -i predicted_flow_%d.jpg -loop -1 flow.gif | |
| def write_flo(filename, flow): | |
| """ | |
| write optical flow in Middlebury .flo format | |
| :param flow: optical flow map | |
| :param filename: optical flow file path to be saved | |
| :return: None | |
| """ | |
| f = open(filename, 'wb') | |
| magic = np.array([202021.25], dtype=np.float32) | |
| (height, width) = flow.shape[0:2] | |
| w = np.array([width], dtype=np.int32) | |
| h = np.array([height], dtype=np.int32) | |
| magic.tofile(f) | |
| w.tofile(f) | |
| h.tofile(f) | |
| flow.tofile(f) | |
| f.close() | |
| gr.Interface(fn=infer, inputs=[], outputs=gr.Textbox()).launch() |