Update app.py

app.py

@@ -7,7 +7,7 @@ import numpy as np
 import os
 import matplotlib.pyplot as plt
 from io import BytesIO
-import base64
+import tempfile
 
 # Check if CUDA is available, otherwise use CPU
 device = 'cuda' if torch.cuda.is_available() else 'cpu'
@@ -70,7 +70,7 @@ def process_video(video_path, target, progress=gr.Progress()):
                 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.5:
+                    if score.item() >= 0.25:
                         box = [round(i, 2) for i in box.tolist()]
                         object_label = target
                         confidence = round(score.item(), 3)
@@ -100,12 +100,11 @@ def create_heatmap(frame_scores):
     plt.yticks([])
     plt.tight_layout()
     
-    buf = BytesIO()
-    plt.savefig(buf, format='png')
-    buf.seek(0)
+    with tempfile.NamedTemporaryFile(delete=False, suffix='.png') as tmp_file:
+        plt.savefig(tmp_file.name, format='png')
     plt.close()
     
-    return base64.b64encode(buf.getvalue()).decode('utf-8')
+    return tmp_file.name
 
 def load_sample_frame(video_path):
     cap = cv2.VideoCapture(video_path)
@@ -120,7 +119,7 @@ def load_sample_frame(video_path):
 
 def gradio_app():
     with gr.Blocks() as app:
-        gr.Markdown("# Video Object Detection with Owlv2")
+        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")
@@ -138,8 +137,8 @@ def gradio_app():
         def process_and_update(video, target):
             frames, scores, error = process_video(video, target, progress_bar)
             if frames is not None:
-                heatmap = create_heatmap(scores)
-                return frames, scores, frames[0], heatmap, error, gr.Slider(maximum=len(frames) - 1, value=0)
+                heatmap_path = create_heatmap(scores)
+                return frames, scores, frames[0], heatmap_path, error, gr.Slider(maximum=len(frames) - 1, value=0)
             return None, None, None, None, error, gr.Slider(maximum=100, value=0)
 
         def update_frame(frame_index, frames):