Update app.py

app.py

@@ -33,7 +33,7 @@ def yolov9_inference(model_id, img_path=None, vid_path=None, tracking_algorithm
         img.save(img_path)
         input_path = img_path
 
-        output_path, df, frame_counts_df = run(weights=model_id, imgsz=(image_size,image_size), conf_thres=conf_threshold, iou_thres=iou_threshold, source=input_path, device='0', hide_conf= True, hide_labels = True)
+        output_path, df, frame_counts_df = run(weights=model_id, imgsz=(image_size,image_size), conf_thres=conf_threshold, iou_thres=iou_threshold, source=input_path, device='cpu', hide_conf= True)
     elif vid_path is not None:
         vid_name = 'output.mp4'
 
@@ -70,12 +70,12 @@ def yolov9_inference(model_id, img_path=None, vid_path=None, tracking_algorithm
         out.release()
         input_path = vid_name
         if tracking_algorithm == 'deep_sort':
-            output_path, df, frame_counts_df = run_deepsort(weights=model_id, imgsz=(image_size,image_size), conf_thres=conf_threshold, iou_thres=iou_threshold, source=input_path, device='0', draw_trails=True)
+            output_path, df, frame_counts_df = run_deepsort(weights=model_id, imgsz=(image_size,image_size), conf_thres=conf_threshold, iou_thres=iou_threshold, source=input_path, device='cpu', draw_trails=True)
         elif tracking_algorithm == 'strong_sort':
-            device_strongsort = torch.device('cuda:0')
-            output_path, df, frame_counts_df = run_strongsort(yolo_weights=model_id, imgsz=(image_size,image_size), conf_thres=conf_threshold, iou_thres=iou_threshold, source=input_path, device=device_strongsort, strong_sort_weights = "osnet_x0_25_msmt17.pt", hide_conf= True, hide_labels = True)
+            device_strongsort = torch.device('cpu')
+            output_path, df, frame_counts_df = run_strongsort(yolo_weights=model_id, imgsz=(image_size,image_size), conf_thres=conf_threshold, iou_thres=iou_threshold, source=input_path, device=device_strongsort, strong_sort_weights = "osnet_x0_25_msmt17.pt", hide_conf= True)
         else: 
-            output_path, df, frame_counts_df =  run(weights=model_id, imgsz=(image_size,image_size), conf_thres=conf_threshold, iou_thres=iou_threshold, source=input_path, device='0', hide_conf= True, hide_labels = True)
+            output_path, df, frame_counts_df =  run(weights=model_id, imgsz=(image_size,image_size), conf_thres=conf_threshold, iou_thres=iou_threshold, source=input_path, device='cpu', hide_conf= True)
         # Assuming output_path is the path to the output file
     _, output_extension = os.path.splitext(output_path)
     palette = {"Bus": "red", "Bike": "blue", "Car": "green", "Pedestrian": "yellow", "Truck": "purple"}
@@ -86,16 +86,25 @@ def yolov9_inference(model_id, img_path=None, vid_path=None, tracking_algorithm
         fig, ax = plt.subplots(figsize=(10, 6))
         #for label in labels:
             #df_label = frame_counts_df[frame_counts_df['label'] == label]
-        sns.barplot(ax = ax, data = df, x = 'label', y = 'count', palette=palette, hue = 'label', legend = False)
+        
+        sns.barplot(ax=ax, data=df, x='label', y='count', palette=palette, hue='label')
 
         # Customizations
-        ax.set_title('Number of Objects', fontsize=20)
-        ax.set_xlabel('Object Class', fontsize=15)
-        ax.set_ylabel('Object Count', fontsize=15)
-        ax.tick_params(axis='x', rotation=45)  # Rotate x-axis labels for better readability
+        ax.set_title('Count of Labels', fontsize=20, pad=20)  # Increase padding for the title
+        ax.set_xlabel('Label', fontsize=16)  # Increase font size
+        ax.set_ylabel('Count', fontsize=16)  # Increase font size
+        ax.tick_params(axis='x', rotation=45, labelsize=12)  # Increase label size and rotate x-axis labels for better readability
+        ax.tick_params(axis='y', labelsize=12)  # Increase label size for y-axis
         sns.despine()  # Remove the top and right spines from plot
-        #ax.legend()
-        ax.grid(True)
+
+        # Add grid but make it lighter and put it behind bars
+        ax.grid(True, linestyle=':', linewidth=0.6, color='gray', alpha=0.6)
+        ax.set_axisbelow(True)
+
+        # Add a legend with a smaller font size
+        ax.legend(fontsize=10)
+
+        plt.tight_layout()  # Ensure the entire plot fits into the figure area
         #ax.set_facecolor('#D3D3D3')
     elif output_extension.lower() in vid_extensions:
         output_video = output_path  # Load the video file here
@@ -104,14 +113,25 @@ def yolov9_inference(model_id, img_path=None, vid_path=None, tracking_algorithm
         fig, ax = plt.subplots(figsize=(10, 6))
         #for label in labels:
             #df_label = frame_counts_df[frame_counts_df['label'] == label]
-        sns.lineplot(ax = ax, data = frame_counts_df[::30],  x = 'frame', y = 'count', hue = 'label', palette = palette)
-
-        ax.set_xlabel('Second')
-        ax.set_ylabel('Object Count')
-        ax.set_title('Number of Objects over Seconds')
-        ax.legend()
-        ax.grid(True)
-        ax.set_facecolor('#D3D3D3')
+        sns.lineplot(ax = ax, data = frame_counts_df[::4],  x = 'frame', y = 'count', hue = 'label', palette=palette, linewidth=2.5)
+
+        ax.set_title('Count of Labels over Frames', fontsize=20, pad=20)  # Increase padding for the title
+        ax.set_xlabel('Frame', fontsize=16)  # Increase font size
+        ax.set_ylabel('Count', fontsize=16)  # Increase font size
+        ax.tick_params(axis='x', labelsize=12)  # Increase label size for x-axis
+        ax.tick_params(axis='y', labelsize=12)  # Increase label size for y-axis
+
+        # Add grid but make it lighter and put it behind bars
+        ax.grid(True, linestyle=':', linewidth=0.6, color='gray', alpha=0.6)
+        ax.set_axisbelow(True)
+
+        # Change the background color to a lighter shade
+        ax.set_facecolor('#F0F0F0')
+
+        # Add a legend with a smaller font size
+        ax.legend(fontsize=10)
+
+        plt.tight_layout()  # Ensure the entire
     return output_image, output_video, fig
 
 def app():
@@ -167,6 +187,7 @@ def app():
             outputs=[output_image, output_video, fig],
         )
 
+
 gradio_app = gr.Blocks()
 with gradio_app:
     gr.HTML(
@@ -188,4 +209,3 @@ with gradio_app:
             app()
 
 gradio_app.launch(debug=True)
-