add app.py

.gitattributes

@@ -33,3 +33,6 @@ saved_model/**/* 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
+football.mp4 filter=lfs diff=lfs merge=lfs -text
+cat.mp4 filter=lfs diff=lfs merge=lfs -text
+safari2.mp4 filter=lfs diff=lfs merge=lfs -text

README.md

@@ -1,6 +1,6 @@
 ---
 title: Omdet Turbo Open Vocabulary
-emoji: 💻
+emoji: 📹
 colorFrom: red
 colorTo: blue
 sdk: gradio
@@ -8,6 +8,8 @@ sdk_version: 4.42.0
 app_file: app.py
 pinned: false
 license: apache-2.0
+short_description: Video captioning/open-vocabulary/zero-shot
+
 ---
 
 Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -0,0 +1,180 @@
+import os
+import time
+
+import cv2
+import gradio as gr
+import numpy as np
+import spaces
+import supervision as sv
+import torch
+from PIL import Image
+from tqdm import tqdm
+
+from transformers import AutoModelForZeroShotObjectDetection, AutoProcessor
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+processor = AutoProcessor.from_pretrained("omdet-turbo-tiny-timm")
+model = AutoModelForZeroShotObjectDetection.from_pretrained("omdet-turbo-tiny-timm").to(
+    device
+)
+
+css = """
+#warning {background-color: #FFCCCB}
+.feedback textarea {font-size: 24px !important}
+"""
+
+BOUNDING_BOX_ANNOTATOR = sv.BoundingBoxAnnotator()
+MASK_ANNOTATOR = sv.MaskAnnotator()
+LABEL_ANNOTATOR = sv.LabelAnnotator()
+
+
+def calculate_end_frame_index(source_video_path):
+    video_info = sv.VideoInfo.from_video_path(source_video_path)
+    return min(video_info.total_frames, video_info.fps * 2)
+
+
+def annotate_image(input_image, detections, labels) -> np.ndarray:
+    output_image = MASK_ANNOTATOR.annotate(input_image, detections)
+    output_image = BOUNDING_BOX_ANNOTATOR.annotate(output_image, detections)
+    output_image = LABEL_ANNOTATOR.annotate(output_image, detections, labels=labels)
+    return output_image
+
+
+def resize_to_max_side(frame: np.ndarray, max_side: int = 640):
+    h, w = frame.shape[:2]
+    if h > w:
+        new_h, new_w = max_side, int(w * max_side / h)
+    else:
+        new_h, new_w = int(h * max_side / w), max_side
+
+    return cv2.resize(frame, (new_w, new_h))
+
+
+@spaces.GPU
+def process_video(
+    input_video,
+    confidence_threshold,
+    classes,
+    max_side,
+    progress=gr.Progress(track_tqdm=True),
+):
+    classes = classes.strip(" ").split(",")
+    video_info = sv.VideoInfo.from_video_path(input_video)
+    total = calculate_end_frame_index(input_video)
+    frame_generator = sv.get_video_frames_generator(source_path=input_video, end=total)
+
+    result_file_name = "output.mp4"
+    result_file_path = os.path.join(os.getcwd(), result_file_name)
+    all_fps = []
+    with sv.VideoSink(result_file_path, video_info=video_info) as sink:
+        for _ in tqdm(range(total), desc="Processing video.."):
+            frame = next(frame_generator)
+            results, fps = query(
+                frame, classes, confidence_threshold, max_side=max_side
+            )
+            all_fps.append(fps)
+            detections = []
+
+            detections = sv.Detections(
+                xyxy=results[0]["boxes"].cpu().detach().numpy(),
+                confidence=results[0]["scores"].cpu().detach().numpy(),
+                class_id=np.array(
+                    [
+                        classes.index(results_class)
+                        for results_class in results[0]["classes"]
+                    ]
+                ),
+                data={"class_name": results[0]["classes"]},
+            )
+            frame = annotate_image(
+                input_image=frame,
+                detections=detections,
+                labels=results[0]["classes"],
+            )
+            sink.write_frame(frame)
+
+    avg_fps = np.mean(all_fps)
+    return result_file_path, gr.Markdown(
+        f'<h3 style="text-align: center;">Model inference FPS: {avg_fps:.2f}</h3>',
+        visible=True,
+    )
+
+
+def query(frame, classes, confidence_threshold, max_side=360):
+    frame_resized = resize_to_max_side(frame, max_side=max_side)
+    image = Image.fromarray(cv2.cvtColor(frame_resized, cv2.COLOR_BGR2RGB))
+    inputs = processor(images=image, text=classes, return_tensors="pt").to(device)
+    with torch.no_grad():
+        start = time.time()
+        outputs = model(**inputs)
+        fps = 1 / (time.time() - start)
+    target_sizes = [frame.shape[:2]]
+
+    results = processor.post_process_grounded_object_detection(
+        outputs=outputs,
+        classes=classes,
+        score_threshold=confidence_threshold,
+        target_sizes=target_sizes,
+    )
+    return results, fps
+
+
+with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
+    gr.Markdown("## Real Time Open Vocabulary Object Detection with Omdet-Turbo")
+    gr.Markdown(
+        "This is a demo for open vocabulary object detection using OmDet-Turbo.  \\"
+        "It runs on ZeroGPU which captures GPU every first time you infer. This combined with video processing time means that the demo inference time is slower than the model's actual inference time.  \\"
+        "The actual model inference FPS is displayed under the processed video after inference."
+    )
+    gr.Markdown(
+        "Simply upload a video, and write the objects you want to detect! You can also play with confidence threshold, image size, or try the examples below. 👇"
+    )
+
+    with gr.Row():
+        with gr.Column():
+            input_video = gr.Video(label="Input Video")
+            submit = gr.Button()
+        with gr.Column():
+            output_video = gr.Video(label="Output Video")
+            actual_fps = gr.Markdown("", visible=False)
+    with gr.Row():
+        classes = gr.Textbox(
+            "person, cat, dog",
+            label="Objects to detect. Change this as you like!",
+            elem_classes="feedback",
+            scale=3,
+        )
+        conf = gr.Slider(
+            label="Confidence Threshold",
+            minimum=0.1,
+            maximum=1.0,
+            value=0.2,
+            step=0.05,
+        )
+        max_side = gr.Slider(
+            label="Image Size",
+            minimum=240,
+            maximum=1080,
+            value=640,
+            step=10,
+        )
+    example = gr.Examples(
+        fn=process_video,
+        examples=[
+            ["./football.mp4", 0.3, "person, ball, shoe", 640],
+            ["./cat.mp4", 0.2, "cat", 640],
+            ["./safari2.mp4", 0.3, "elephant, giraffe, springbok, zebra", 640],
+        ],
+        inputs=[input_video, conf, classes, max_side],
+        outputs=output_video,
+    )
+
+    submit.click(
+        fn=process_video,
+        inputs=[input_video, conf, classes, max_side],
+        outputs=[output_video, actual_fps],
+    )
+
+if __name__ == "__main__":
+    demo.launch(show_error=True)

cat.mp4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:07539c031a516acecf58b8751f74ba90182efe4c4ad25513038f10564739eadd
+size 810095

football.mp4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:56a85c5c7d5d6e0825f76a71e5e3ee2ce35c8ffbe841ef4bfa544af1089259aa
+size 2855852

requirement.txt

@@ -0,0 +1,5 @@
+torch
+timm
+git+https://github.com/yonigozlan/transformers.git@add-om-det-turbo
+supervision
+spaces

safari2.mp4

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