Update app.py

app.py

@@ -222,6 +222,9 @@ def process_video(confidence_threshold=0.5, selected_classes=None, stream_url=No
         # Perform object tracking with confidence threshold on the resized frame
         results = model.track(resized_frame, persist=True, conf=confidence_threshold)
 
+        # Create an annotator for the original frame
+        annotator = Annotator(frame, line_width=2)
+
         if results[0].boxes.id is not None:
             track_ids = results[0].boxes.id.int().cpu().tolist()
             clss = results[0].boxes.cls.cpu().tolist()
@@ -232,6 +235,11 @@ def process_video(confidence_threshold=0.5, selected_classes=None, stream_url=No
                 if conf >= confidence_threshold and model.names[cls] in selected_classes:
                     # Scale the bounding box back to the original resolution
                     box = box * (original_width / new_width)
+                    box = box.int().tolist()
+
+                    # Draw the bounding box on the original frame
+                    annotator.box_label(box, label=f"{model.names[cls]} {conf:.2f}", color=colors(cls))
+
                     # Check if the object crosses the line
                     if is_object_crossing_line(box, line_params) and t_id not in crossed_objects:
                         crossed_objects[t_id] = True
@@ -240,11 +248,8 @@ def process_video(confidence_threshold=0.5, selected_classes=None, stream_url=No
                         if len(crossed_objects) > max_tracked_objects:
                             crossed_objects.clear()
 
-        # Visualize the results with bounding boxes, masks, and IDs on the original frame
-        annotated_frame = results[0].plot(img=frame)
-
         # Draw the angled line on the original frame
-        draw_angled_line(annotated_frame, line_params, color=(0, 255, 0), thickness=2)
+        draw_angled_line(frame, line_params, color=(0, 255, 0), thickness=2)
 
         # Display the count on the frame with a modern look
         count = len(crossed_objects)
@@ -252,17 +257,17 @@ def process_video(confidence_threshold=0.5, selected_classes=None, stream_url=No
 
         # Calculate the position for the middle of the top
         margin = 10  # Margin from the top
-        x = (annotated_frame.shape[1] - text_width) // 2  # Center-align the text horizontally
+        x = (frame.shape[1] - text_width) // 2  # Center-align the text horizontally
         y = text_height + margin  # Top-align the text
 
         # Draw the black background rectangle
-        cv2.rectangle(annotated_frame, (x - margin, y - text_height - margin), (x + text_width + margin, y + margin), (0, 0, 0), -1)
+        cv2.rectangle(frame, (x - margin, y - text_height - margin), (x + text_width + margin, y + margin), (0, 0, 0), -1)
 
         # Draw the text
-        cv2.putText(annotated_frame, f"COUNT: {count}", (x, y), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
+        cv2.putText(frame, f"COUNT: {count}", (x, y), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
 
         # Yield the annotated frame to Gradio
-        yield annotated_frame, ""
+        yield frame, ""
 
     cap.release()
     logger.info("Stream processing completed.")