Create app.py

app.py

@@ -0,0 +1,414 @@
+import cv2
+import numpy as np
+from collections import deque
+from datetime import datetime
+from ultralytics import YOLO
+import time
+
+class RotatingPadShirtCounter:
+    """
+    Robust shirt counter for rotating pad system.
+    Logic: Count when empty pad ENTERS the ROI (after shirt was removed)
+    """
+    def __init__(self, 
+                 model_path='runs/exp2/weights/best.pt',
+                 roi_center=(320, 240), 
+                 roi_radius=180,
+                 min_conf=0.5,
+                 stability_frames=5):
+        
+        # Load YOLO model
+        print(f"Loading YOLO model from: {model_path}")
+        self.model = YOLO(model_path)
+        self.model_names = self.model.names
+        print(f"Model classes: {self.model_names}")
+        
+        # ROI Configuration
+        self.roi_center = roi_center
+        self.roi_radius = roi_radius
+        self.min_conf = min_conf
+        
+        # State tracking
+        self.current_state = "UNKNOWN"
+        self.prev_state = "UNKNOWN"
+        self.state_buffer = deque(maxlen=stability_frames)
+        self.stability_frames = stability_frames
+        
+        # Counting logic
+        self.shirt_count = 0
+        
+        # Prevent double counting
+        self.last_count_time = time.time()
+        self.min_time_between_counts = 3.0
+        
+        # Detection history
+        self.detection_history = deque(maxlen=30)
+        
+        self.pad_away_frames = 0
+        self.min_pad_away_frames = 80
+
+        # Logging
+        self.event_log = []
+        self.debug_mode = True
+        
+    def detect_in_roi(self, frame):
+        """
+        Run YOLO detection and filter by ROI
+        Returns: (has_empty_pad, has_occupied_pad, all_detections)
+        """
+        # Run YOLO
+        results = self.model.predict(frame, conf=self.min_conf, verbose=False)
+        
+        has_empty_pad_in_roi = False
+        has_occupied_pad_in_roi = False
+        all_detections = []
+        
+        # Parse results
+        for result in results:
+            boxes = result.boxes
+            
+            for box in boxes:
+                # Extract data
+                x1, y1, x2, y2 = box.xyxy[0].cpu().numpy()
+                conf = float(box.conf[0].cpu().numpy())
+                class_id = int(box.cls[0].cpu().numpy())
+                class_name = self.model_names[class_id]
+                
+                # Calculate center
+                center_x = (x1 + x2) / 2
+                center_y = (y1 + y2) / 2
+                
+                # Check if in ROI
+                dist = np.sqrt((center_x - self.roi_center[0])**2 + 
+                              (center_y - self.roi_center[1])**2)
+                
+                in_roi = dist < self.roi_radius
+                
+                detection = {
+                    'bbox': [x1, y1, x2, y2],
+                    'center': (center_x, center_y),
+                    'confidence': conf,
+                    'class': class_name,
+                    'in_roi': in_roi
+                }
+                all_detections.append(detection)
+                
+                if in_roi:
+                    if class_name == 'empty_pad':
+                        has_empty_pad_in_roi = True
+                    else:
+                        # Any other detection in ROI means occupied (shirt on pad)
+                        has_occupied_pad_in_roi = True
+        
+        return has_empty_pad_in_roi, has_occupied_pad_in_roi, all_detections
+    
+    def determine_state(self, has_empty, has_occupied):
+        """Determine current state based on detections"""
+        if has_empty:
+            return "EMPTY_IN_ROI"
+        elif has_occupied:
+            return "OCCUPIED_IN_ROI"
+        else:
+            return "PAD_AWAY"
+    
+    def update_state_buffer(self, state):
+        """Add to buffer and return stable state"""
+        self.state_buffer.append(state)
+        
+        if len(self.state_buffer) < self.stability_frames:
+            return self.current_state
+        
+        # Count occurrences
+        state_counts = {}
+        for s in self.state_buffer:
+            state_counts[s] = state_counts.get(s, 0) + 1
+        
+        # Get most common state
+        stable_state = max(state_counts, key=state_counts.get)
+        
+        # Require majority agreement (> 60%)
+        if state_counts[stable_state] >= len(self.state_buffer) * 0.6:
+            return stable_state
+        
+        return self.current_state
+    
+    def should_count(self):
+        """
+        KEY COUNTING LOGIC:
+        Count when worker removes shirt: OCCUPIED_IN_ROI -> EMPTY_IN_ROI
+        But only if previous PAD_AWAY state lasted >= 80 frames
+        """
+        if self.prev_state == "PAD_AWAY" and self.current_state == "OCCUPIED_IN_ROI":
+            time_since_last = time.time() - self.last_count_time
+            if (time_since_last >= self.min_time_between_counts and 
+                self.pad_away_frames >= self.min_pad_away_frames):
+                return True, f"Shirt on pad after PAD_AWAY for {self.pad_away_frames} frames"
+        
+        return False, None
+    
+    def process_frame(self, frame):
+        """Main processing loop"""
+        # Detect
+        has_empty, has_occupied, detections = self.detect_in_roi(frame)
+        
+        # Determine instantaneous state
+        instant_state = self.determine_state(has_empty, has_occupied)
+        
+        # Get stable state
+        stable_state = self.update_state_buffer(instant_state)
+        
+        # Track how long previous state was PAD_AWAY
+        if self.current_state == "PAD_AWAY":
+            self.pad_away_frames += 1
+        else:
+            self.pad_away_frames = 0  # Reset if not PAD_AWAY
+        
+        # Check for state change
+        state_changed = (stable_state != self.current_state)
+        
+        if state_changed:
+            self.prev_state = self.current_state
+            self.current_state = stable_state
+            
+            # Check if we should count
+            should_count, reason = self.should_count()
+            
+            if should_count:
+                self.shirt_count += 1
+                self.last_count_time = time.time()
+                self.log_event("SHIRT_COUNTED", reason)
+                print(f"🎯 SHIRT #{self.shirt_count} COUNTED! - {reason}")
+            else:
+                self.log_event("STATE_CHANGE", f"{self.prev_state} -> {self.current_state}")
+        
+        # Visualize
+        vis_frame = self.draw_visualization(frame, detections, instant_state)
+        
+        return vis_frame
+    
+    def draw_visualization(self, frame, detections, instant_state):
+        """Draw debug information on frame"""
+        vis = frame.copy()
+        
+        # Draw ROI
+        cv2.circle(vis, self.roi_center, self.roi_radius, (0, 255, 255), 3)
+        cv2.circle(vis, self.roi_center, 5, (0, 255, 255), -1)
+        
+        # Draw all detections
+        for det in detections:
+            x1, y1, x2, y2 = map(int, det['bbox'])
+            conf = det['confidence']
+            cls = det['class']
+            in_roi = det['in_roi']
+            
+            # Color coding
+            if cls == 'empty_pad':
+                color = (0, 255, 0)  # Green
+            else:
+                color = (0, 0, 255)  # Red
+            
+            thickness = 3 if in_roi else 2
+            cv2.rectangle(vis, (x1, y1), (x2, y2), color, thickness)
+            
+            # Label
+            label = f"{cls} {conf:.2f}"
+            if in_roi:
+                label += " [ROI]"
+            cv2.putText(vis, label, (x1, y1-10),
+                       cv2.FONT_HERSHEY_SIMPLEX, 0.6, color, 2)
+        
+        # Status panel
+        panel_height = 180
+        panel = np.zeros((panel_height, vis.shape[1], 3), dtype=np.uint8)
+        
+        # Count (BIG and prominent)
+        cv2.putText(panel, f"SHIRTS COUNTED: {self.shirt_count}", (20, 50),
+                   cv2.FONT_HERSHEY_SIMPLEX, 1.5, (0, 255, 0), 3)
+        
+        # Current state
+        state_color = {
+            "EMPTY_IN_ROI": (0, 255, 0),
+            "OCCUPIED_IN_ROI": (0, 165, 255),
+            "PAD_AWAY": (255, 0, 0),
+            "UNKNOWN": (128, 128, 128)
+        }.get(self.current_state, (255, 255, 255))
+        
+        cv2.putText(panel, f"State: {self.current_state}", (20, 90),
+                   cv2.FONT_HERSHEY_SIMPLEX, 0.8, state_color, 2)
+        
+        # Instant vs Stable
+        cv2.putText(panel, f"Instant: {instant_state}", (20, 120),
+                   cv2.FONT_HERSHEY_SIMPLEX, 0.6, (200, 200, 200), 1)
+        
+        # Buffer visualization
+        buffer_str = ''.join([
+            'E' if s == "EMPTY_IN_ROI" else 
+            'O' if s == "OCCUPIED_IN_ROI" else 
+            'A' if s == "PAD_AWAY" else '?'
+            for s in self.state_buffer
+        ])
+        cv2.putText(panel, f"Buffer: [{buffer_str}]", (20, 150),
+                   cv2.FONT_HERSHEY_SIMPLEX, 0.6, (180, 180, 180), 1)
+        
+        # Combine
+        vis = np.vstack([panel, vis])
+        
+        return vis
+    
+    def log_event(self, event_type, details):
+        """Log events for debugging"""
+        self.event_log.append({
+            'timestamp': datetime.now().strftime('%H:%M:%S.%f')[:-3],
+            'event': event_type,
+            'details': details,
+            'count': self.shirt_count,
+            'state': self.current_state
+        })
+    
+    def get_stats(self):
+        """Get statistics"""
+        return {
+            'total_shirts': self.shirt_count,
+            'current_state': self.current_state,
+            'events': self.event_log
+        }
+
+
+# ============================================================================
+# MAIN CONFIGURATION - EDIT THESE VALUES
+# ============================================================================
+
+# INPUT/OUTPUT FILES
+INPUT_VIDEO = "videos/sdcard_0_20251013125904.mp4"           # Your input video path
+OUTPUT_VIDEO = "videos/output_counted.mp4"       # Where to save processed video
+MODEL_PATH = "runs/exp2/weights/best.pt"  # Your YOLO model path
+
+# ROI SETTINGS (Region of Interest where pad appears)
+ROI_CENTER = None        # None = auto-detect (video center), or tuple like (640, 360)
+ROI_RADIUS = 180         # Radius in pixels
+
+# DETECTION SETTINGS
+MIN_CONFIDENCE = 0.98     # Minimum YOLO confidence (0.0 to 1.0)
+STABILITY_FRAMES = 15     # Frames needed to confirm state change
+
+# ============================================================================
+
+
+def process_video():
+    """
+    Process video file and save output with detections and counting
+    """
+    print("="*80)
+    print("ROTATING PAD SHIRT COUNTER - VIDEO PROCESSOR")
+    print("="*80)
+    
+    # Open input video
+    cap = cv2.VideoCapture(INPUT_VIDEO)
+    if not cap.isOpened():
+        print(f"❌ Error: Cannot open video file: {INPUT_VIDEO}")
+        return
+    
+    # Get video properties
+    fps = int(cap.get(cv2.CAP_PROP_FPS))
+    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+    
+    print(f"✓ Input Video: {INPUT_VIDEO}")
+    print(f"  Resolution: {width}x{height}")
+    print(f"  FPS: {fps}")
+    print(f"  Total Frames: {total_frames}")
+    
+    # Auto-calculate ROI center if not provided
+    roi_center = ROI_CENTER if ROI_CENTER else (width // 2, height // 2)
+    
+    print(f"  ROI Center: {roi_center}, Radius: {ROI_RADIUS}")
+    
+    # Initialize counter
+    counter = RotatingPadShirtCounter(
+        model_path=MODEL_PATH,
+        roi_center=roi_center,
+        roi_radius=ROI_RADIUS,
+        min_conf=MIN_CONFIDENCE,
+        stability_frames=STABILITY_FRAMES
+    )
+    
+    # Prepare output video writer
+    output_height = height + 180  # Add panel height
+    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
+    out = cv2.VideoWriter(OUTPUT_VIDEO, fourcc, fps, (width, output_height))
+    
+    if not out.isOpened():
+        print(f"❌ Error: Cannot create output video: {OUTPUT_VIDEO}")
+        cap.release()
+        return
+    
+    print(f"✓ Output Video: {OUTPUT_VIDEO}")
+    print(f"  Output Resolution: {width}x{output_height}")
+    print("-"*80)
+    print("Processing video...")
+    
+    frame_count = 0
+    start_time = time.time()
+    
+    try:
+        while True:
+            ret, frame = cap.read()
+            if not ret:
+                break
+            
+            frame_count += 1
+            
+            # Process frame
+            vis_frame = counter.process_frame(frame)
+            
+            # Add frame number and progress
+            progress = (frame_count / total_frames) * 100
+            cv2.putText(vis_frame, f"Frame: {frame_count}/{total_frames} ({progress:.1f}%)", 
+                       (width - 350, 30),
+                       cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 0), 2)
+            
+            # Write frame
+            out.write(vis_frame)
+            
+            # Progress indicator every 30 frames
+            if frame_count % 30 == 0:
+                elapsed = time.time() - start_time
+                fps_processing = frame_count / elapsed
+                eta_seconds = (total_frames - frame_count) / fps_processing if fps_processing > 0 else 0
+                print(f"Progress: {frame_count}/{total_frames} frames "
+                      f"({progress:.1f}%) | "
+                      f"Shirts: {counter.shirt_count} | "
+                      f"ETA: {eta_seconds:.0f}s")
+    
+    except KeyboardInterrupt:
+        print("\n⚠ Processing interrupted by user")
+    
+    finally:
+        # Cleanup
+        cap.release()
+        out.release()
+        
+        # Final statistics
+        elapsed = time.time() - start_time
+        stats = counter.get_stats()
+        
+        print("\n" + "="*80)
+        print("PROCESSING COMPLETE")
+        print("="*80)
+        print(f"Total Frames Processed: {frame_count}")
+        print(f"Processing Time: {elapsed:.2f} seconds")
+        print(f"Average FPS: {frame_count/elapsed:.2f}")
+        print(f"\n🎯 TOTAL SHIRTS COUNTED: {stats['total_shirts']}")
+        print(f"Final State: {stats['current_state']}")
+        print("\nEvent Log (Shirt Counts):")
+        for evt in stats['events']:
+            if evt['event'] == 'SHIRT_COUNTED':
+                print(f"  ✓ [{evt['timestamp']}] Shirt #{evt['count']} - {evt['details']}")
+        print("="*80)
+        print(f"✓ Output saved to: {OUTPUT_VIDEO}")
+        print("="*80)
+
+
+if __name__ == "__main__":
+    process_video()