Update app.py

app.py

@@ -6,7 +6,6 @@ from transformers import Owlv2Processor, Owlv2ForObjectDetection
 import numpy as np
 import os
 import matplotlib.pyplot as plt
-from io import BytesIO
 import tempfile
 import shutil
 
@@ -24,7 +23,6 @@ except RuntimeError:
     device = torch.device("cpu")
     model = model.to(device)
 
-
 def process_video(video_path, target, progress=gr.Progress()):
     if video_path is None:
         return None, None, "Error: No video uploaded"
@@ -46,16 +44,7 @@ def process_video(video_path, target, progress=gr.Progress()):
     temp_dir = tempfile.mkdtemp()
     frame_paths = []
 
-    # Try to use GPU with half precision, fall back to CPU if out of memory
-    try:
-        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-        model.to(device).half()  # Convert model to half precision
-    except RuntimeError:
-        print("GPU out of memory, falling back to CPU")
-        device = torch.device("cpu")
-        model.to(device)
-
-    batch_size = 1 
+    batch_size = 4  # Process 4 frames at a time
     batch_frames = []
     batch_indices = []
 
@@ -67,8 +56,8 @@ def process_video(video_path, target, progress=gr.Progress()):
             break
 
         # Resize the frame
-        #img_resized = cv2.resize(img, (1280, 720))
-        pil_img = Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
+        img_resized = cv2.resize(img, (640, 360))
+        pil_img = Image.fromarray(cv2.cvtColor(img_resized, cv2.COLOR_BGR2RGB))
         
         batch_frames.append(pil_img)
         batch_indices.append(i)
@@ -87,42 +76,36 @@ def process_video(video_path, target, progress=gr.Progress()):
                 draw = ImageDraw.Draw(pil_img)
                 max_score = 0
 
-                try:
-                    font = ImageFont.truetype("arial.ttf", 20)
-                except IOError:
-                    font = ImageFont.load_default()
-
                 boxes, scores, labels = result["boxes"], result["scores"], result["labels"]
 
-for box, score, label in zip(boxes, scores, labels):
-    if score.item() >= 0.5:
-        box = [round(i, 2) for i in box.tolist()]
-        object_label = target
-        confidence = round(score.item(), 3)
-        annotation = f"{object_label}: {confidence}"
+                for box, score, label in zip(boxes, scores, labels):
+                    if score.item() >= 0.5:
+                        box = [round(i, 2) for i in box.tolist()]
+                        object_label = target
+                        confidence = round(score.item(), 3)
+                        annotation = f"{object_label}: {confidence}"
 
-        # Increase line width for the bounding box
-        draw.rectangle(box, outline="red", width=4)
+                        # Increase line width for the bounding box
+                        draw.rectangle(box, outline="red", width=4)
 
-        # Increase font size and change color to red
-        font_size = 30  # Increased from 20
-        try:
-            font = ImageFont.truetype("arial.ttf", font_size)
-        except IOError:
-            font = ImageFont.load_default()
+                        # Increase font size and change color to red
+                        font_size = 30
+                        try:
+                            font = ImageFont.truetype("arial.ttf", font_size)
+                        except IOError:
+                            font = ImageFont.load_default()
 
-        text_position = (box[0], box[1] - font_size - 5)
-        
-        # Add a semi-transparent background for better text visibility
-        text_bbox = draw.textbbox(text_position, annotation, font=font)
-        draw.rectangle(text_bbox, fill=(0, 0, 0, 128))
+                        text_position = (box[0], box[1] - font_size - 5)
+                        
+                        # Add a semi-transparent background for better text visibility
+                        text_bbox = draw.textbbox(text_position, annotation, font=font)
+                        draw.rectangle(text_bbox, fill=(0, 0, 0, 128))
 
-        # Draw text in red
-        draw.text(text_position, annotation, fill="red", font=font)
+                        # Draw text in red
+                        draw.text(text_position, annotation, fill="red", font=font)
 
-        max_score = max(max_score, confidence)
+                        max_score = max(max_score, confidence)
 
-                # Save frame to disk
                 frame_path = os.path.join(temp_dir, f"frame_{batch_indices[idx]:04d}.png")
                 pil_img.save(frame_path)
                 frame_paths.append(frame_path)
@@ -146,13 +129,11 @@ def create_heatmap(frame_scores, current_frame):
     plt.xlabel('Frame', fontsize=12)
     plt.yticks([])
 
-    # Add more frame numbers on x-axis
     num_frames = len(frame_scores)
-    step = max(1, num_frames // 20)  # Show at most 20 frame numbers
+    step = max(1, num_frames // 20)
     frame_numbers = range(0, num_frames, step)
     plt.xticks(frame_numbers, [str(i) for i in frame_numbers], rotation=45, ha='right')
 
-    # Add vertical line for current frame
     plt.axvline(x=current_frame, color='blue', linestyle='--', linewidth=2)
 
     plt.tight_layout()
@@ -234,7 +215,7 @@ def gradio_app():
 
 if __name__ == "__main__":
     app = gradio_app()
-    app.launch(share=True)
+    app.launch()
 
     # Cleanup temporary files
     def cleanup():