Update app.py

app.py

@@ -16,7 +16,7 @@ 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)
 
-@spaces.GPU(duration=200)
+@spaces.GPU(duration=120)
 def process_video(video_path, target, progress=gr.Progress()):
     if video_path is None:
         return None, None, "Error: No video uploaded"
@@ -36,9 +36,6 @@ def process_video(video_path, target, progress=gr.Progress()):
 
     processed_frames = []
     frame_scores = []
-    batch_size = 2
-    batch_frames = []
-    batch_times = []
 
     for time in progress.tqdm(np.arange(0, video_duration, frame_duration)):
         frame_number = int(time * original_fps)
@@ -48,47 +45,39 @@ def process_video(video_path, target, progress=gr.Progress()):
             break
 
         pil_img = Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
-        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", 40)
-                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=2)
-                        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 = []
+        
+        # Process single image
+        inputs = processor(text=[target], images=pil_img, return_tensors="pt", padding=True).to(device)
+        outputs = model(**inputs)
+
+        target_sizes = torch.Tensor([pil_img.size[::-1]])
+        results = processor.post_process_object_detection(outputs, target_sizes=target_sizes)
+        
+        draw = ImageDraw.Draw(pil_img)
+        max_score = 0
+
+        try:
+            font = ImageFont.truetype("arial.ttf", 40)
+        except IOError:
+            font = ImageFont.load_default()
+
+        boxes, scores, labels = results[0]["boxes"], results[0]["scores"], results[0]["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=2)
+                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(pil_img))
+        frame_scores.append(max_score)
 
     cap.release()
     return processed_frames, frame_scores, None