Duplicate from fffiloni/RAFT

Co-authored-by: Sylvain Filoni <fffiloni@users.noreply.huggingface.co>

.gitattributes

@@ -0,0 +1,34 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,13 @@
+---
+title: RAFT Optical Flow
+emoji: 😻
+colorFrom: green
+colorTo: blue
+sdk: gradio
+sdk_version: 3.19.1
+app_file: app.py
+pinned: false
+duplicated_from: fffiloni/RAFT
+---
+
+Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -0,0 +1,178 @@
+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 cv2
+import numpy as np
+import os
+import sys
+import torch
+from PIL import Image
+import matplotlib.pyplot as plt
+import torchvision.transforms.functional as F
+from torchvision.io import read_video, read_image, ImageReadMode
+from torchvision.models.optical_flow import Raft_Large_Weights
+from torchvision.models.optical_flow import raft_large
+from torchvision.io import write_jpeg
+import torchvision.transforms as T
+
+import tempfile
+from pathlib import Path
+from urllib.request import urlretrieve
+
+from scipy.interpolate import LinearNDInterpolator
+from imageio import imread, imwrite
+
+
+
+def write_flo(flow, filename):
+    """
+    Write optical flow in Middlebury .flo format
+    
+    :param flow: optical flow map
+    :param filename: optical flow file path to be saved
+    :return: None
+    
+    from https://github.com/liruoteng/OpticalFlowToolkit/
+    
+    """
+    # forcing conversion to float32 precision
+    flow = flow.cpu().data.numpy()
+    flow = flow.astype(np.float32)
+    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()
+
+
+    
+def infer(frameA, frameB):
+    #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("./spacex.mp4"), output_format="TCHW")
+    #print(f"FRAME BEFORE stack: {frames[100]}")
+    
+    
+    input_frame_1 = read_image(str(frameA), ImageReadMode.UNCHANGED)
+    print(f"FRAME 1: {input_frame_1}")
+    input_frame_2 = read_image(str(frameB), ImageReadMode.UNCHANGED)
+    print(f"FRAME 1: {input_frame_2}")
+    
+    #img1_batch = torch.stack([frames[0]])
+    #img2_batch = torch.stack([frames[1]])
+
+    img1_batch = torch.stack([input_frame_1])
+    img2_batch = torch.stack([input_frame_2])
+    
+    print(f"FRAME AFTER stack: {img1_batch}")
+    
+    weights = Raft_Large_Weights.DEFAULT
+    transforms = weights.transforms()
+
+
+    def preprocess(img1_batch, img2_batch):
+        img1_batch = F.resize(img1_batch, size=[520, 960])
+        img2_batch = F.resize(img2_batch, size=[520, 960])
+        return transforms(img1_batch, img2_batch)
+
+
+    img1_batch, img2_batch = preprocess(img1_batch, img2_batch)
+
+    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.
+
+    
+
+# 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_batch.to(device), img2_batch.to(device))
+    print(f"list_of_flows type = {type(list_of_flows)}")
+    print(f"list_of_flows 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"predicted_flows dtype = {predicted_flows.dtype}")
+    print(f"predicted_flows shape = {predicted_flows.shape} = (N, 2, H, W)")
+    print(f"predicted_flows min = {predicted_flows.min()}, predicted_flows 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)
+
+    predicted_flow = list_of_flows[-1][0]
+    print(f"predicted flow dtype = {predicted_flow.dtype}")
+    print(f"predicted flow shape = {predicted_flow.shape}")
+    
+    flow_img = flow_to_image(predicted_flow).to("cpu")
+    write_jpeg(flow_img, f"predicted_flow.jpg")
+    
+    flo_file = write_flo(predicted_flow, "flofile.flo")
+    
+    return "predicted_flow.jpg", ["flofile.flo"]
+
+description="<p style='text-align:center'>PyTorch way to Generate optical flow image & .flo file from 2 consecutive frames with RAFT model</p>" 
+
+gr.Interface(fn=infer, inputs=[gr.Image(source="upload", type="filepath", label="frame 1"), gr.Image(source="upload", type="filepath", label="frame 2")], outputs=[gr.Image(label="flow image"), gr.Files(label="flow file")], title="RAFT Optical Flow", description=description).launch()

requirements.txt

@@ -0,0 +1,13 @@
+--extra-index-url https://download.pytorch.org/whl/cu113
+av
+torch
+torchvision
+pathlib
+matplotlib
+opencv-contrib-python
+scipy
+imageio
+git+https://github.com/huggingface/diffusers.git 
+transformers
+git+https://github.com/huggingface/accelerate
+xformers==0.0.16