Update app.py

app.py

@@ -94,18 +94,33 @@ def process_video(video_path, target, progress=gr.Progress()):
 
                 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 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()
+
+        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.rectangle(box, outline="red", width=2)
-                        text_position = (box[0], box[1] - 20)
-                        draw.text(text_position, annotation, fill="white", 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")
@@ -125,17 +140,17 @@ def process_video(video_path, target, progress=gr.Progress()):
     return frame_paths, frame_scores, None
 
 def create_heatmap(frame_scores, current_frame):
-    plt.figure(figsize=(12, 3))
+    plt.figure(figsize=(16, 4))
     plt.imshow([frame_scores], cmap='hot_r', aspect='auto')
-    plt.title('Object Detection Heatmap')
-    plt.xlabel('Frame')
+    plt.title('Object Detection Heatmap', fontsize=14)
+    plt.xlabel('Frame', fontsize=12)
     plt.yticks([])
 
     # Add more frame numbers on x-axis
     num_frames = len(frame_scores)
-    step = max(1, num_frames // 10)  # Show at most 10 frame numbers
+    step = max(1, num_frames // 20)  # Show at most 20 frame numbers
     frame_numbers = range(0, num_frames, step)
-    plt.xticks(frame_numbers, [str(i) for i in frame_numbers])
+    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)