Update app.py

app.py

@@ -15,42 +15,6 @@ 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", padding=True).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", 30)
-    except IOError:
-        font = ImageFont.load_default()
-
-    i = 0
-    text = texts[i]
-    boxes, scores, labels = results[i]["boxes"], results[i]["scores"], results[i]["labels"]
-
-    max_score = 0
-    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=4)
-            text_position = (box[0], box[1] - 30)
-            draw.text(text_position, annotation, fill="white", font=font)
-
-            max_score = max(max_score, confidence)
-
-    return image, max_score
-
 def process_video(video_path, target, progress=gr.Progress()):
     if video_path is None:
         return None, None, "Error: No video uploaded"
@@ -70,6 +34,9 @@ def process_video(video_path, target, progress=gr.Progress()):
 
     processed_frames = []
     frame_scores = []
+    batch_size = 32
+    batch_frames = []
+    batch_times = []
 
     for time in progress.tqdm(np.arange(0, video_duration, frame_duration)):
         frame_number = int(time * original_fps)
@@ -79,9 +46,47 @@ def process_video(video_path, target, progress=gr.Progress()):
             break
 
         pil_img = Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
-        annotated_img, max_score = detect_objects_in_frame(pil_img, target)
-        processed_frames.append(np.array(annotated_img))
-        frame_scores.append(max_score)
+        batch_frames.append(pil_img)
+        batch_times.append(time)
+
+        if len(batch_frames) == batch_size or time + frame_duration >= video_duration:
+            # Process the batch
+            texts = [[target]] * len(batch_frames)
+            inputs = processor(text=texts, images=batch_frames, return_tensors="pt", padding=True).to(device)
+            outputs = model(**inputs)
+
+            for i, (image, batch_time) in enumerate(zip(batch_frames, batch_times)):
+                target_sizes = torch.Tensor([image.size[::-1]])
+                results = processor.post_process_object_detection(outputs, target_sizes=target_sizes)
+                
+                draw = ImageDraw.Draw(image)
+                max_score = 0
+
+                try:
+                    font = ImageFont.truetype("arial.ttf", 30)
+                except IOError:
+                    font = ImageFont.load_default()
+
+                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 = target
+                        confidence = round(score.item(), 3)
+                        annotation = f"{object_label}: {confidence}"
+
+                        draw.rectangle(box, outline="red", width=4)
+                        text_position = (box[0], box[1] - 30)
+                        draw.text(text_position, annotation, fill="white", font=font)
+
+                        max_score = max(max_score, confidence)
+
+                processed_frames.append(np.array(image))
+                frame_scores.append(max_score)
+
+            batch_frames = []
+            batch_times = []
 
     cap.release()
     return processed_frames, frame_scores, None
@@ -115,7 +120,7 @@ def load_sample_frame(video_path):
 
 def gradio_app():
     with gr.Blocks() as app:
-        gr.Markdown("# Video Object Detection with Owlv2 (3 FPS)")
+        gr.Markdown("# Video Object Detection with Owlv2 (3 FPS, Batch Size 32)")
 
         video_input = gr.Video(label="Upload Video")
         target_input = gr.Textbox(label="Target Object", value="Elephant")