Update app.py

app.py

@@ -34,7 +34,7 @@ def process_video(video_path, target, progress=gr.Progress()):
 
     processed_frames = []
     frame_scores = []
-    batch_size = 1
+    batch_size = 2
     batch_frames = []
     batch_times = []
 
@@ -63,7 +63,7 @@ def process_video(video_path, target, progress=gr.Progress()):
                 max_score = 0
 
                 try:
-                    font = ImageFont.truetype("arial.ttf", 30)
+                    font = ImageFont.truetype("arial.ttf", 40)
                 except IOError:
                     font = ImageFont.load_default()
 
@@ -76,7 +76,7 @@ def process_video(video_path, target, progress=gr.Progress()):
                         confidence = round(score.item(), 3)
                         annotation = f"{object_label}: {confidence}"
 
-                        draw.rectangle(box, outline="red", width=4)
+                        draw.rectangle(box, outline="red", width=2)
                         text_position = (box[0], box[1] - 30)
                         draw.text(text_position, annotation, fill="white", font=font)
 
@@ -91,17 +91,29 @@ def process_video(video_path, target, progress=gr.Progress()):
     cap.release()
     return processed_frames, frame_scores, None
 
-def create_heatmap(frame_scores):
-    plt.figure(figsize=(10, 2))
-    plt.imshow([frame_scores], cmap='hot', aspect='auto')
-    plt.colorbar(label='Confidence')
+def create_heatmap(frame_scores, current_frame):
+    plt.figure(figsize=(12, 3))
+    plt.imshow([frame_scores], cmap='hot_r', aspect='auto')  # 'hot_r' for reversed hot colormap
+    cbar = plt.colorbar(label='Confidence')
+    cbar.ax.yaxis.set_ticks_position('left')
+    cbar.ax.yaxis.set_label_position('left')
     plt.title('Object Detection Heatmap')
     plt.xlabel('Frame')
     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
+    frame_numbers = range(0, num_frames, step)
+    plt.xticks(frame_numbers, [str(i) for i in frame_numbers])
+
+    # Add vertical line for current frame
+    plt.axvline(x=current_frame, color='blue', linestyle='--', linewidth=2)
+
     plt.tight_layout()
     
     with tempfile.NamedTemporaryFile(delete=False, suffix='.png') as tmp_file:
-        plt.savefig(tmp_file.name, format='png')
+        plt.savefig(tmp_file.name, format='png', dpi=400, bbox_inches='tight')
     plt.close()
     
     return tmp_file.name
@@ -124,10 +136,10 @@ def gradio_app():
         video_input = gr.Video(label="Upload Video")
         target_input = gr.Textbox(label="Target Object", value="Elephant")
         frame_slider = gr.Slider(minimum=0, maximum=100, step=1, label="Frame", value=0)
-        output_image = gr.Image(label="Processed Frame")
         heatmap_output = gr.Image(label="Detection Heatmap")
+        output_image = gr.Image(label="Processed Frame")
         error_output = gr.Textbox(label="Error Messages", visible=False)
-        sample_video_frame = gr.Image(value=load_sample_frame("Drone Video of African Wildlife Wild Botswan.mp4"), label="Sample Video Frame")
+        sample_video_frame = gr.Image(value=load_sample_frame("Drone Video of African Wildlife Wild Botswan.mp4"), label="Drone Video of African Wildlife Wild Botswan by wildimagesonline.com - Sample Video Frame")
         use_sample_button = gr.Button("Use Sample Video")
         progress_bar = gr.Progress()
 
@@ -137,22 +149,23 @@ def gradio_app():
         def process_and_update(video, target):
             frames, scores, error = process_video(video, target, progress_bar)
             if frames is not None:
-                heatmap_path = create_heatmap(scores)
+                heatmap_path = create_heatmap(scores, 0)  # Initial heatmap with current frame at 0
                 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):
+        def update_frame_and_heatmap(frame_index, frames, scores):
             if frames and 0 <= frame_index < len(frames):
-                return frames[frame_index]
-            return None
+                heatmap_path = create_heatmap(scores, frame_index)
+                return frames[frame_index], heatmap_path
+            return None, None
 
         video_input.upload(process_and_update, 
                            inputs=[video_input, target_input], 
                            outputs=[processed_frames, frame_scores, output_image, heatmap_output, error_output, frame_slider])
 
-        frame_slider.change(update_frame, 
-                            inputs=[frame_slider, processed_frames], 
-                            outputs=[output_image])
+        frame_slider.change(update_frame_and_heatmap, 
+                            inputs=[frame_slider, processed_frames, frame_scores], 
+                            outputs=[output_image, heatmap_output])
 
         def use_sample_video():
             sample_video_path = "Drone Video of African Wildlife Wild Botswan.mp4"
@@ -162,6 +175,19 @@ def gradio_app():
                                 inputs=None, 
                                 outputs=[processed_frames, frame_scores, output_image, heatmap_output, error_output, frame_slider])
 
+        # Rearrange the layout
+        video_input.render()
+        target_input.render()
+        with gr.Row():
+            with gr.Column(scale=2):
+                output_image.render()
+            with gr.Column(scale=1):
+                sample_video_frame.render()
+                use_sample_button.render()
+        frame_slider.render()
+        heatmap_output.render()
+        error_output.render()
+
     return app
 
 if __name__ == "__main__":