Update app.py

app.py

@@ -21,39 +21,11 @@ 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
 
 
-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
@@ -64,29 +36,15 @@ def infer():
     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), output_format="TCHW")
 
-    img1= [frames[100]
-    img2 = [frames[101]
+    
+    frames, _, _ = read_video(str(video_path), output_format="TCHW")
 
-#########################
-# The RAFT model accepts RGB images. We first get the frames from
-# :func:`~torchvision.io.read_video` and resize them to ensure their
-# dimensions are divisible by 8. Then we use the transforms bundled into the
-# weights in order to preprocess the input and rescale its values to the
-# required ``[-1, 1]`` interval.
+    img1= frames[100]
+    img2 = frames[101]
 
-    from torchvision.models.optical_flow import Raft_Large_Weights
+    
 
     weights = Raft_Large_Weights.DEFAULT
     transforms = weights.transforms()
@@ -112,7 +70,7 @@ def infer():
 # 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
+    
 
 # If you can, run this example on a GPU, it will be a lot faster.
     device = "cuda" if torch.cuda.is_available() else "cpu"