Create app.py

app.py

@@ -0,0 +1,129 @@
+import gradio as gr
+import cv2
+from PIL import Image, ImageDraw, ImageFont
+import torch
+from transformers import Owlv2Processor, Owlv2ForObjectDetection
+import numpy as np
+import os
+
+# Check if CUDA is available, otherwise use CPU
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
+
+processor = Owlv2Processor.from_pretrained("google/owlv2-base-patch16")
+model = Owlv2ForObjectDetection.from_pretrained("google/owlv2-base-patch16").to(device)
+
+def detect_objects_in_frame(image, target):
+    draw = ImageDraw.Draw(image)
+    texts = [[target]]
+    inputs = processor(text=texts, images=image, return_tensors="pt").to(device)
+    outputs = model(**inputs)
+
+    target_sizes = torch.Tensor([image.size[::-1]])
+    results = processor.post_process_object_detection(outputs=outputs, threshold=0.1, target_sizes=target_sizes)
+
+    color_map = {target: "red"}
+
+    try:
+        font = ImageFont.truetype("arial.ttf", 15)
+    except IOError:
+        font = ImageFont.load_default()
+
+    i = 0
+    text = texts[i]
+    boxes, scores, labels = results[i]["boxes"], results[i]["scores"], results[i]["labels"]
+
+    for box, score, label in zip(boxes, scores, labels):
+        if score.item() >= 0.25:
+            box = [round(i, 2) for i in box.tolist()]
+            object_label = text[label]
+            confidence = round(score.item(), 3)
+            annotation = f"{object_label}: {confidence}"
+
+            draw.rectangle(box, outline=color_map.get(object_label, "red"), width=2)
+            text_position = (box[0], box[1] - 10)
+            draw.text(text_position, annotation, fill="white", font=font)
+
+    return image
+
+def process_video(video_path, target, progress=gr.Progress()):
+    if video_path is None:
+        return None, "Error: No video uploaded"
+
+    if not os.path.exists(video_path):
+        return None, f"Error: Video file not found at {video_path}"
+
+    cap = cv2.VideoCapture(video_path)
+    if not cap.isOpened():
+        return None, f"Error: Unable to open video file at {video_path}"
+
+    frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+    original_fps = int(cap.get(cv2.CAP_PROP_FPS))
+    original_duration = frame_count / original_fps
+
+    output_path = "output_video.mp4"
+    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
+    out = cv2.VideoWriter(output_path, fourcc, original_fps, (int(cap.get(3)), int(cap.get(4))))
+
+    for frame in progress.tqdm(range(frame_count)):
+        ret, img = cap.read()
+        if not ret:
+            break
+
+        pil_img = Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
+        annotated_img = detect_objects_in_frame(pil_img, target)
+        annotated_frame = cv2.cvtColor(np.array(annotated_img), cv2.COLOR_RGB2BGR)
+        out.write(annotated_frame)
+
+    cap.release()
+    out.release()
+
+    return output_path, None
+
+def load_sample_frame(video_path):
+    cap = cv2.VideoCapture(video_path)
+    if not cap.isOpened():
+        return None
+    ret, frame = cap.read()
+    cap.release()
+    if not ret:
+        return None
+    frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+    return frame_rgb
+
+def gradio_app():
+    with gr.Blocks() as app:
+        gr.Markdown("# Video Object Detection with Owlv2")
+
+        video_input = gr.Video(label="Upload Video")
+        target_input = gr.Textbox(label="Target Object")
+        output_video = gr.Video(label="Output Video")
+        error_output = gr.Textbox(label="Error Messages", visible=False)
+        sample_video_frame = gr.Image(value=load_sample_frame("IL_Dancing_Sample.mp4"), label="Sample Video Frame")
+        use_sample_button = gr.Button("Use Sample Video")
+
+        video_path = gr.State(None)
+        def process_and_update(video, target):
+            output_video_path, error = process_video(video, target)
+            if error:
+                error_output.visible = True
+            else:
+                error_output.visible = False
+            return output_video_path, error
+
+        video_input.upload(process_and_update, 
+                           inputs=[video_input, target_input], 
+                           outputs=[output_video, error_output])
+
+        def use_sample_video():
+            sample_video_path = "IL_Dancing_Sample.mp4"
+            return process_and_update(sample_video_path, "person")
+
+        use_sample_button.click(use_sample_video, 
+                                inputs=None, 
+                                outputs=[output_video, error_output])
+
+    return app
+
+if __name__ == "__main__":
+    app = gradio_app()
+    app.launch(share=True)