Update

.gitattributes

@@ -34,3 +34,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
 *.jpg filter=lfs diff=lfs merge=lfs -text
+*.mp4 filter=lfs diff=lfs merge=lfs -text

app.py

@@ -7,16 +7,22 @@ https://colab.research.google.com/drive/1e8fcby5rhKZWcr9LSN8mNbQ0TU4Dxxpo
 """
 
 import pathlib
+import tempfile
 
+import cv2
 import gradio as gr
+import numpy as np
 import PIL.Image
 import spaces
 import supervision as sv
 import torch
+import tqdm
 from transformers import AutoProcessor, RTDetrForObjectDetection, VitPoseForPoseEstimation
 
 DESCRIPTION = "# ViTPose"
 
+MAX_NUM_FRAMES = 300
+
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 person_detector_name = "PekingU/rtdetr_r50vd_coco_o365"
@@ -30,7 +36,7 @@ pose_model = VitPoseForPoseEstimation.from_pretrained(pose_model_name, device_ma
 
 @spaces.GPU
 @torch.inference_mode()
-def run(image: PIL.Image.Image) -> tuple[PIL.Image.Image, list[dict]]:
+def process_image(image: PIL.Image.Image) -> tuple[PIL.Image.Image, list[dict]]:
     inputs = person_image_processor(images=image, return_tensors="pt").to(device)
     outputs = person_model(**inputs)
     results = person_image_processor.post_process_object_detection(
@@ -101,26 +107,82 @@ def run(image: PIL.Image.Image) -> tuple[PIL.Image.Image, list[dict]]:
     return vertex_annotator.annotate(scene=annotated_frame, key_points=keypoints), human_readable_results
 
 
-paths = sorted(pathlib.Path("images").glob("*.jpg"))
+def process_video(
+    video_path: str,
+    progress: gr.Progress = gr.Progress(track_tqdm=True),  # noqa: ARG001, B008
+) -> str:
+    cap = cv2.VideoCapture(video_path)
+
+    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+    fps = cap.get(cv2.CAP_PROP_FPS)
+    num_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+
+    fourcc = cv2.VideoWriter_fourcc(*"mp4v")
+    with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as out_file:
+        writer = cv2.VideoWriter(out_file.name, fourcc, fps, (width, height))
+        for _ in tqdm.auto.tqdm(range(min(MAX_NUM_FRAMES, num_frames))):
+            ok, frame = cap.read()
+            if not ok:
+                break
+            rgb_frame = frame[:, :, ::-1]
+            annotated_frame, _ = process_image(PIL.Image.fromarray(rgb_frame))
+            writer.write(np.asarray(annotated_frame)[:, :, ::-1])
+        writer.release()
+    cap.release()
+    return out_file.name
+
+
+process_video.zerogpu = True  # type: ignore
 
 
 with gr.Blocks(css_paths="style.css") as demo:
     gr.Markdown(DESCRIPTION)
-    with gr.Row():
-        with gr.Column():
-            input_image = gr.Image(label="Input Image", type="pil")
-            run_button = gr.Button()
-        with gr.Column():
-            output_image = gr.Image(label="Output Image")
-            output_json = gr.JSON(label="Output JSON")
-
-    gr.Examples(examples=paths, inputs=input_image, outputs=[output_image, output_json], fn=run)
-
-    run_button.click(
-        fn=run,
-        inputs=input_image,
-        outputs=[output_image, output_json],
-    )
+
+    with gr.Tabs():
+        with gr.Tab("Image"):
+            with gr.Row():
+                with gr.Column():
+                    input_image = gr.Image(label="Input Image", type="pil")
+                    run_button_image = gr.Button()
+                with gr.Column():
+                    output_image = gr.Image(label="Output Image")
+                    output_json = gr.JSON(label="Output JSON")
+            gr.Examples(
+                examples=sorted(pathlib.Path("images").glob("*.jpg")),
+                inputs=input_image,
+                outputs=[output_image, output_json],
+                fn=process_image,
+            )
+
+            run_button_image.click(
+                fn=process_image,
+                inputs=input_image,
+                outputs=[output_image, output_json],
+            )
+
+        with gr.Tab("Video"):
+            gr.Markdown(f"The input video will be truncated to {MAX_NUM_FRAMES} frames.")
+
+            with gr.Row():
+                with gr.Column():
+                    input_video = gr.Video(label="Input Video")
+                    run_button_video = gr.Button()
+                with gr.Column():
+                    output_video = gr.Video(label="Output Video")
+
+            gr.Examples(
+                examples=sorted(pathlib.Path("videos").glob("*.mp4")),
+                inputs=input_video,
+                outputs=output_video,
+                fn=process_video,
+                cache_examples=False,
+            )
+            run_button_video.click(
+                fn=process_video,
+                inputs=input_video,
+                outputs=output_video,
+            )
 
 
 if __name__ == "__main__":

pyproject.toml

@@ -8,6 +8,7 @@ dependencies = [
     "accelerate>=1.2.1",
     "gradio>=5.11.0",
     "hf-transfer>=0.1.9",
+    "opencv-python-headless>=4.10.0.84",
     "setuptools>=75.8.0",
     "spaces>=0.32.0",
     "supervision>=0.25.1",

requirements.txt

@@ -108,6 +108,7 @@ numpy==2.2.1
     #   gradio
     #   matplotlib
     #   opencv-python
+    #   opencv-python-headless
     #   pandas
     #   scipy
     #   supervision
@@ -145,6 +146,8 @@ nvidia-nvtx-cu12==12.1.105
     # via torch
 opencv-python==4.10.0.84
     # via supervision
+opencv-python-headless==4.10.0.84
+    # via vitpose-transformers (pyproject.toml)
 orjson==3.10.14
     # via gradio
 packaging==24.2

uv.lock

@@ -937,6 +937,23 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/ec/6c/fab8113424af5049f85717e8e527ca3773299a3c6b02506e66436e19874f/opencv_python-4.10.0.84-cp37-abi3-win_amd64.whl", hash = "sha256:32dbbd94c26f611dc5cc6979e6b7aa1f55a64d6b463cc1dcd3c95505a63e48fe", size = 38842521 },
 ]
 
+[[package]]
+name = "opencv-python-headless"
+version = "4.10.0.84"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "numpy" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/2f/7e/d20f68a5f1487adf19d74378d349932a386b1ece3be9be9915e5986db468/opencv-python-headless-4.10.0.84.tar.gz", hash = "sha256:f2017c6101d7c2ef8d7bc3b414c37ff7f54d64413a1847d89970b6b7069b4e1a", size = 95117755 }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/1c/9b/583c8d9259f6fc19413f83fd18dd8e6cbc8eefb0b4dc6da52dd151fe3272/opencv_python_headless-4.10.0.84-cp37-abi3-macosx_11_0_arm64.whl", hash = "sha256:a4f4bcb07d8f8a7704d9c8564c224c8b064c63f430e95b61ac0bffaa374d330e", size = 54835657 },
+    { url = "https://files.pythonhosted.org/packages/c0/7b/b4c67f5dad7a9a61c47f7a39e4050e8a4628bd64b3c3daaeb755d759f928/opencv_python_headless-4.10.0.84-cp37-abi3-macosx_12_0_x86_64.whl", hash = "sha256:5ae454ebac0eb0a0b932e3406370aaf4212e6a3fdb5038cc86c7aea15a6851da", size = 56475470 },
+    { url = "https://files.pythonhosted.org/packages/91/61/f838ce2046f3ec3591ea59ea3549085e399525d3b4558c4ed60b55ed88c0/opencv_python_headless-4.10.0.84-cp37-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:46071015ff9ab40fccd8a163da0ee14ce9846349f06c6c8c0f2870856ffa45db", size = 29329705 },
+    { url = "https://files.pythonhosted.org/packages/d1/09/248f86a404567303cdf120e4a301f389b68e3b18e5c0cc428de327da609c/opencv_python_headless-4.10.0.84-cp37-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:377d08a7e48a1405b5e84afcbe4798464ce7ee17081c1c23619c8b398ff18295", size = 49858781 },
+    { url = "https://files.pythonhosted.org/packages/30/c0/66f88d58500e990a9a0a5c06f98862edf1d0a3a430781218a8c193948438/opencv_python_headless-4.10.0.84-cp37-abi3-win32.whl", hash = "sha256:9092404b65458ed87ce932f613ffbb1106ed2c843577501e5768912360fc50ec", size = 28675298 },
+    { url = "https://files.pythonhosted.org/packages/26/d0/22f68eb23eea053a31655960f133c0be9726c6a881547e6e9e7e2a946c4f/opencv_python_headless-4.10.0.84-cp37-abi3-win_amd64.whl", hash = "sha256:afcf28bd1209dd58810d33defb622b325d3cbe49dcd7a43a902982c33e5fad05", size = 38754031 },
+]
+
 [[package]]
 name = "orjson"
 version = "3.10.14"
@@ -1815,6 +1832,7 @@ dependencies = [
     { name = "accelerate" },
     { name = "gradio" },
     { name = "hf-transfer" },
+    { name = "opencv-python-headless" },
     { name = "setuptools" },
     { name = "spaces" },
     { name = "supervision" },
@@ -1827,6 +1845,7 @@ requires-dist = [
     { name = "accelerate", specifier = ">=1.2.1" },
     { name = "gradio", specifier = ">=5.11.0" },
     { name = "hf-transfer", specifier = ">=0.1.9" },
+    { name = "opencv-python-headless", specifier = ">=4.10.0.84" },
     { name = "setuptools", specifier = ">=75.8.0" },
     { name = "spaces", specifier = ">=0.32.0" },
     { name = "supervision", specifier = ">=0.25.1" },

videos/README.md

@@ -0,0 +1,6 @@
+These videos are from the following public domain:
+
+- https://www.pexels.com/video/young-guy-doing-break-dance-on-the-street-5362370/
+- https://www.pexels.com/video/a-woman-dancing-at-home-6003986/
+- https://www.pexels.com/video/long-haired-man-dancing-in-a-library-6344381/
+- https://www.pexels.com/video/a-female-model-dancing-around-6815069/

videos/pexels-allan-mas-5362370.mp4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:747f9c2f9d19e4955603e1a13b69663187882d4c6a8fbcad18ddbd04ee792d4d
+size 1972564

videos/pexels-artem-podrez-6003986.mp4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1044083afc06aa6f956838c7fcd582c9cfd59ea3a994adc8a0f5889ffca4d9c8
+size 2494082

videos/pexels-c-technical-6344381.mp4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7763476045f4683d53d751fb8befaf637c0101a0693e72f5b582e6aa5ac63cac
+size 3967587

videos/pexels-roman-odintsov-6815069.mp4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:44045b239c0f523bfeedc5871019ae9f67525fcf65ba46d7ca4516994e6b2f57
+size 2617714