Repository: pytorch/vision

Issue:
Sphinx-gallery example for RAFT  model
Now that we have a `flow_to_image()` util https://github.com/pytorch/vision/issues/5091, we can write a sphinx-gallery example to illustrate how to use the RAFT model.

Ideally, the example would look like something like this:

- Briefly explain what optical flow is: predicting movement from 2 consecutive frames
- load the model with pre-trained weights
- load a video (from our `assets` folder), or download it. This is TBD. If we host the video on our GitHub repo we need to make sure it's free of use **and** super lightweight, so it's unlikely. One potential option might be to host the video on our S3 bucket where we host the pre-trained weighst?
- read frames from this video with `read_video()`. Ideally we would be using `VideoReader()`, but it is only available to users who build torchvision from source, so it may not be usable by most. Note: we don't want to rely on external dependencies to read the video
- pass consecutive frames to the model, get predictions, and convert the predicted flows to images via `flow_to_image()`
- plot a few of the predictions for illustration
- save the predicted images to a folder and tell the user they can create a video with e.g. `ffmpeg -f image2 -framerate 30 -i frame_%d.jpg -loop -1 my_cool_gif.gif`

A good starting point is https://github.com/pytorch/vision/issues/5091#issuecomment-992677348

Some caveats worth noting:
- we need the example to be reasonably fast to run, so that means the image resolution should not be too big. Ideally it should take < 1 min on the CI (which does not have a GPU!)
- If loading a video is impossible (or hard), we can fallback to just using pre-selected pairs of images. The example won't be as fancy, but still interesting, and we can expand later.
- loading the weights from `torchvision.models.optical_flow` will give pre-trained weights from Chairs + Things. These may not be the most accurate weights, depending on the type of video that we apply the model to. The most accurate weights would probably be e.g. `Raft_Large_Weights.C_T_SKHT_V2` or `C_T_SKHT_K_V2`, but these are only available through the prototype API in `torchvision.prototype.models.optical_flow`. Whether we can already use this API in a sphinx-gallery example is TBD: perhaps @datumbox can share his thoughts here? In any case it's safe to start writing the eaxmple with the basic pre-trained weights from `torchvision.models.optical_flow`.


As discussed offline @oke-aditya is interested in this so I'll assign it to you :)


Discussion:
> If we host the video on our GitHub repo we need to make sure it's free of use and super lightweight

Using one of the assets included in our repo sounds like a good choice.

> The most accurate weights would probably be e.g. Raft_Large_Weights.C_T_SKHT_V2 or C_T_SKHT_K_V2, but these are only available through the prototype API in torchvision.prototype.models.optical_flow. 

The multi-weight API is still in prototype because the community has been providing feedback (https://github.com/pytorch/vision/issues/5088) and making changes up until last week. That's why it's probably hard to have it moved on main TorchVision area in this release. This means that you can't depend on it from examples because it will cause issues on the release (unless you change it on the release branch). Another approach could be to change the default weights of your non-prototype model builder (`pretrained=True`) to use the ones from `Raft_Large_Weights.C_T_SKHT_V2`. The change is only a few lines of code and since they haven't been officially released, it's non-BC breaking.

Resolving patch:
diff --git a/gallery/plot_optical_flow.py b/gallery/plot_optical_flow.py
new file mode 100644
--- /dev/null
+++ b/gallery/plot_optical_flow.py
@@ -0,0 +1,198 @@
+"""
+=====================================================
+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
+import torchvision.transforms as T
+
+
+plt.rcParams["savefig.bbox"] = "tight"
+# sphinx_gallery_thumbnail_number = 2
+
+
+def plot(imgs, **imshow_kwargs):
+    if not isinstance(imgs[0], list):
+        # Make a 2d grid even if there's just 1 row
+        imgs = [imgs]
+
+    num_rows = len(imgs)
+    num_cols = len(imgs[0])
+    _, axs = plt.subplots(nrows=num_rows, ncols=num_cols, squeeze=False)
+    for row_idx, row in enumerate(imgs):
+        for col_idx, img in enumerate(row):
+            ax = axs[row_idx, col_idx]
+            img = F.to_pil_image(img.to("cpu"))
+            ax.imshow(np.asarray(img), **imshow_kwargs)
+            ax.set(xticklabels=[], yticklabels=[], xticks=[], yticks=[])
+
+    plt.tight_layout()
+
+###################################
+# Reading Videos Using Torchvision
+# --------------------------------
+# We will first read a video using :func:`~torchvision.io.read_video`.
+# Alternatively one can use the new :class:`~torchvision.io.VideoReader` API (if
+# torchvision is built from source).
+# The video we will use here is free of use from `pexels.com
+# <https://www.pexels.com/video/a-man-playing-a-game-of-basketball-5192157/>`_,
+# credits go to `Pavel Danilyuk <https://www.pexels.com/@pavel-danilyuk>`_.
+
+
+import tempfile
+from pathlib import Path
+from urllib.request import urlretrieve
+
+
+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)
+
+#########################
+# :func:`~torchvision.io.read_video` returns the video frames, audio frames and
+# the metadata associated with the video. In our case, we only need the video
+# frames.
+#
+# Here we will just make 2 predictions between 2 pre-selected pairs of frames,
+# namely frames (100, 101) and (150, 151). Each of these pairs corresponds to a
+# single model input.
+
+from torchvision.io import read_video
+frames, _, _ = read_video(str(video_path))
+frames = frames.permute(0, 3, 1, 2)  # (N, H, W, C) -> (N, C, H, W)
+
+img1_batch = torch.stack([frames[100], frames[150]])
+img2_batch = torch.stack([frames[101], frames[151]])
+
+plot(img1_batch)
+
+#########################
+# The RAFT model that we will use accepts RGB float images with pixel values in
+# [-1, 1]. The frames we got from :func:`~torchvision.io.read_video` are int
+# images with values in [0, 255], so we will have to pre-process them. We also
+# reduce the image sizes for the example to run faster. Image dimension must be
+# divisible by 8.
+
+
+def preprocess(batch):
+    transforms = T.Compose(
+        [
+            T.ConvertImageDtype(torch.float32),
+            T.Normalize(mean=0.5, std=0.5),  # map [0, 1] into [-1, 1]
+            T.Resize(size=(520, 960)),
+        ]
+    )
+    batch = transforms(batch)
+    return batch
+
+
+# If you can, run this example on a GPU, it will be a lot faster.
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+img1_batch = preprocess(img1_batch).to(device)
+img2_batch = preprocess(img2_batch).to(device)
+
+print(f"shape = {img1_batch.shape}, dtype = {img1_batch.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.
+
+from torchvision.models.optical_flow import raft_large
+
+model = raft_large(pretrained=True, progress=False).to(device)
+model = model.eval()
+
+list_of_flows = model(img1_batch.to(device), img2_batch.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)
+
+# 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
+# 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 = preprocess(img1[None]).to(device)
+#     img2 = preprocess(img2[None]).to(device)
+
+#     list_of_flows = model(img1_batch, img2_batch)
+#     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
diff --git a/torchvision/models/optical_flow/raft.py b/torchvision/models/optical_flow/raft.py
--- a/torchvision/models/optical_flow/raft.py
+++ b/torchvision/models/optical_flow/raft.py
@@ -586,6 +586,8 @@ def raft_large(*, pretrained=False, progress=True, **kwargs):
     """RAFT model from
     `RAFT: Recurrent All Pairs Field Transforms for Optical Flow <https://arxiv.org/abs/2003.12039>`_.
 
+    Please see the example below for a tutorial on how to use this model.
+
     Args:
         pretrained (bool): Whether to use weights that have been pre-trained on
             :class:`~torchvsion.datasets.FlyingChairs` + :class:`~torchvsion.datasets.FlyingThings3D`
@@ -637,6 +639,8 @@ def raft_small(*, pretrained=False, progress=True, **kwargs):
     """RAFT "small" model from
     `RAFT: Recurrent All Pairs Field Transforms for Optical Flow <https://arxiv.org/abs/2003.12039>`_.
 
+    Please see the example below for a tutorial on how to use this model.
+
     Args:
         pretrained (bool): Whether to use weights that have been pre-trained on
             :class:`~torchvsion.datasets.FlyingChairs` + :class:`~torchvsion.datasets.FlyingThings3D`.


Test patch: