Update app.py

app.py

@@ -159,9 +159,11 @@ def convert_video_to_mp4(input_path, output_path):
 def detect_objects_in_video(video_path):
     temp_output_path = "/tmp/output_video.mp4"
     temp_frames_dir = tempfile.mkdtemp()
-    all_class_count = {}  # To store cumulative counts for all frames
-    nestle_total = 0  # Total Nestlé count
+    
+    all_class_count = {}  # To store cumulative counts for all classes (from previous frames)
+    last_frame_objects = {}  # To store the objects detected in the last frame
     frame_count = 0
+    count_text = ""  # Initialize count_text properly
 
     try:
         # Convert video to MP4 if necessary
@@ -190,34 +192,38 @@ def detect_objects_in_video(video_path):
             frame_path = os.path.join(temp_frames_dir, f"frame_{frame_count}.jpg")
             cv2.imwrite(frame_path, frame)
 
-            # Process predictions for frame
+            # Process predictions for the current frame
             predictions = yolo_model.predict(frame_path, confidence=50, overlap=80).json()
 
-            # Update class count for this frame
-            frame_class_count = {}
+            # Get current frame object counts
+            current_frame_objects = {}
             for prediction in predictions['predictions']:
                 class_name = prediction['class']
-                frame_class_count[class_name] = frame_class_count.get(class_name, 0) + 1
-                cv2.rectangle(frame, (int(prediction['x'] - prediction['width']/2), 
-                                      int(prediction['y'] - prediction['height']/2)),
-                              (int(prediction['x'] + prediction['width']/2), 
-                               int(prediction['y'] + prediction['height']/2)),
-                              (0, 255, 0), 2)
-                cv2.putText(frame, class_name, (int(prediction['x'] - prediction['width']/2),
-                                                int(prediction['y'] - prediction['height']/2 - 10)),
-                            cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
-
-            # Update cumulative count for all frames
-            for class_name, count in frame_class_count.items():
+                current_frame_objects[class_name] = current_frame_objects.get(class_name, 0) + 1
+
+                # Draw bounding boxes on the frame
+                x, y, w, h = prediction['x'], prediction['y'], prediction['width'], prediction['height']
+                cv2.rectangle(frame, (int(x-w/2), int(y-h/2)), (int(x+w/2), int(y+h/2)), (0,255,0), 2)
+                cv2.putText(frame, class_name, (int(x-w/2), int(y-h/2-10)), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0,255,0), 2)
+
+            # Update all class counts: Add objects detected in this frame and remove objects not detected in this frame
+            for class_name, count in current_frame_objects.items():
+                # Add to all_class_count
                 all_class_count[class_name] = all_class_count.get(class_name, 0) + count
 
-            # Update total Nestlé products count
-            nestle_total = sum(all_class_count.values())
+            # Now check for objects that were in last frame but not in the current frame
+            for class_name in list(all_class_count.keys()):
+                # If the object was not detected in the current frame, reduce its count
+                if class_name not in current_frame_objects:
+                    all_class_count[class_name] -= 1
+                    if all_class_count[class_name] <= 0:
+                        del all_class_count[class_name]  # Remove from all_class_count if count drops to 0
 
-            # Create a vertical layout for counts (dynamically updated)
+            # Update count_text with current object count
+            count_text = ""
             for class_name, count in all_class_count.items():
                 count_text += f"{class_name}: {count}\n"
-            count_text += f"\nTotal Product Nestlé: {nestle_total}"
+            count_text += f"\nTotal Products: {sum(all_class_count.values())}"
 
             # Overlay the counts text onto the frame
             y_offset = 20