Skip to content

Ikomia-hub/infer_raft_optical_flow

Repository files navigation

Algorithm icon

infer_raft_optical_flow


Stars Website GitHub
Discord community

Run RAFT optical flow algorithm.

Estimate per-pixel motion between two consecutive frames with a RAFT model which is a composition of CNN and RNN. Models are trained with the Sintel dataset

Example image

🚀 Use with Ikomia API

1. Install Ikomia API

We strongly recommend using a virtual environment. If you're not sure where to start, we offer a tutorial here.

pip install ikomia

2. Create your workflow

from ikomia.core import IODataType
from ikomia.dataprocess import CImageIO
from ikomia.dataprocess.workflow import Workflow
from ikomia.utils.displayIO import display
import cv2

# Init your workflow
wf = Workflow()

# Add RAFT optical flow algorithm
optical_flow = wf.add_task(name="infer_raft_optical_flow", auto_connect=True)

stream = cv2.VideoCapture(0)
while True:
    # Read image from stream
    ret, frame = stream.read()

    # Test if streaming is OK
    if not ret:
        continue

    # Run algorithm on current frame
    # RAFT algorithm need at least 2 frames to give results
    optical_flow.set_input(CImageIO(IODataType.IMAGE, frame), 0)
    optical_flow.run()

    # Get and display results
    image_out = optical_flow.get_output(0)
    if image_out.is_data_available():
        img_res = (image_out.get_image()*255).astype('uint8')
        img_res = cv2.cvtColor(img_res, cv2.COLOR_BGR2RGB)
        display(img_res, title="RAFT", viewer="opencv")

    # Press 'q' to quit the streaming process
    if cv2.waitKey(1) & 0xFF == ord('q'):
        break

# After the loop release the stream object
stream.release()
# Destroy all windows
cv2.destroyAllWindows()

☀️ Use with Ikomia Studio

Ikomia Studio offers a friendly UI with the same features as the API.

  • If you haven't started using Ikomia Studio yet, download and install it from this page.

  • For additional guidance on getting started with Ikomia Studio, check out this blog post.

📝 Set algorithm parameters

# Init your workflow
wf = Workflow()

# Add RAFT optical flow algorithm
optical_flow = wf.add_task(name="infer_raft_optical_flow", auto_connect=True)

optical_flow.set_parameters({
    "small": "True",
    "cuda": "True",
})
  • small (bool, default=True): True to use small model (faster), False to use large model (slower, better quality).
  • cuda (bool, default=True): CUDA acceleration if True, run on CPU otherwise.

🔍 Explore algorithm outputs

Every algorithm produces specific outputs, yet they can be explored them the same way using the Ikomia API. For a more in-depth understanding of managing algorithm outputs, please refer to the documentation.

from ikomia.core import IODataType
from ikomia.dataprocess import CImageIO
from ikomia.dataprocess.workflow import Workflow
import cv2

# Init your workflow
wf = Workflow()

# Add RAFT optical flow algorithm
optical_flow = wf.add_task(name="infer_raft_optical_flow", auto_connect=True)

stream = cv2.VideoCapture(0)
while True:
    # Read image from stream
    ret, frame = stream.read()

    # Test if streaming is OK
    if not ret:
        continue

    # Run algorithm on current frame
    # RAFT algorithm need at least 2 frames to give results
    optical_flow.set_input(CImageIO(IODataType.IMAGE, frame), 0)
    optical_flow.run()

    # Iterate over outputs
    for output in optical_flow.get_outputs():
        # Print information
        print(output)
        # Export it to JSON
        output.to_json()

    # Press 'q' to quit the streaming process
    if cv2.waitKey(1) & 0xFF == ord('q'):
        break

# After the loop release the stream object
stream.release()
# Destroy all windows
cv2.destroyAllWindows()

RAFT algorithm generates 1 output:

  • Optical flow image (CImageIO)