Update app.py

app.py

@@ -1,9 +1,12 @@
+import gradio as gr
 import cv2
 import numpy as np
 from collections import deque
 from datetime import datetime
 from ultralytics import YOLO
 import time
+import tempfile
+import os
 
 class RotatingPadShirtCounter:
     """
@@ -11,7 +14,7 @@ class RotatingPadShirtCounter:
     Logic: Count when empty pad ENTERS the ROI (after shirt was removed)
     """
     def __init__(self, 
-                 model_path='runs/exp2/weights/best.pt',
+                 model_path='best.pt',
                  roi_center=(320, 240), 
                  roi_radius=180,
                  min_conf=0.5,
@@ -52,33 +55,25 @@ class RotatingPadShirtCounter:
         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
+        """Run YOLO detection and filter by ROI"""
         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)
                 
@@ -97,7 +92,6 @@ class RotatingPadShirtCounter:
                     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
@@ -118,26 +112,19 @@ class RotatingPadShirtCounter:
         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
-        """
+        """KEY COUNTING LOGIC"""
         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 
@@ -148,29 +135,21 @@ class RotatingPadShirtCounter:
     
     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
+            self.pad_away_frames = 0
         
-        # 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:
@@ -181,51 +160,38 @@ class RotatingPadShirtCounter:
             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
-            
+            color = (0, 255, 0) if cls == 'empty_pad' else (0, 0, 255)
             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),
@@ -236,11 +202,9 @@ class RotatingPadShirtCounter:
         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 
@@ -250,9 +214,7 @@ class RotatingPadShirtCounter:
         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):
@@ -274,82 +236,51 @@ class RotatingPadShirtCounter:
         }
 
 
-# ============================================================================
-# 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)
+def process_video(video_path, roi_radius, min_confidence, stability_frames, progress=gr.Progress()):
+    """Process uploaded video"""
+    if video_path is None:
+        return None, "⚠️ Please upload a video first!"
     
-    # Open input video
-    cap = cv2.VideoCapture(INPUT_VIDEO)
+    progress(0, desc="Opening video...")
+    
+    cap = cv2.VideoCapture(video_path)
     if not cap.isOpened():
-        print(f"❌ Error: Cannot open video file: {INPUT_VIDEO}")
-        return
+        return None, "❌ Error: Cannot open video file"
     
-    # 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)
+    roi_center = (width // 2, height // 2)
     
-    print(f"  ROI Center: {roi_center}, Radius: {ROI_RADIUS}")
+    progress(0.1, desc="Loading model...")
     
-    # Initialize counter
     counter = RotatingPadShirtCounter(
-        model_path=MODEL_PATH,
+        model_path='best.pt',
         roi_center=roi_center,
-        roi_radius=ROI_RADIUS,
-        min_conf=MIN_CONFIDENCE,
-        stability_frames=STABILITY_FRAMES
+        roi_radius=int(roi_radius),
+        min_conf=min_confidence,
+        stability_frames=int(stability_frames)
     )
     
-    # Prepare output video writer
-    output_height = height + 180  # Add panel height
+    output_height = height + 180
+    
+    # Create temporary output file
+    temp_output = tempfile.NamedTemporaryFile(delete=False, suffix='.mp4')
+    output_path = temp_output.name
+    temp_output.close()
+    
     fourcc = cv2.VideoWriter_fourcc(*'mp4v')
-    out = cv2.VideoWriter(OUTPUT_VIDEO, fourcc, fps, (width, output_height))
+    out = cv2.VideoWriter(output_path, fourcc, fps, (width, output_height))
     
     if not out.isOpened():
-        print(f"❌ Error: Cannot create output video: {OUTPUT_VIDEO}")
         cap.release()
-        return
+        return None, "❌ Error: Cannot create output video"
     
-    print(f"✓ Output Video: {OUTPUT_VIDEO}")
-    print(f"  Output Resolution: {width}x{output_height}")
-    print("-"*80)
-    print("Processing video...")
+    progress(0.2, desc="Processing video...")
     
     frame_count = 0
-    start_time = time.time()
     
     try:
         while True:
@@ -358,57 +289,127 @@ def process_video():
                 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}%)", 
+            frame_progress = (frame_count / total_frames) * 100
+            cv2.putText(vis_frame, f"Frame: {frame_count}/{total_frames} ({frame_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")
+                progress(0.2 + (frame_count / total_frames) * 0.75, 
+                        desc=f"Processing: {frame_count}/{total_frames} frames | Shirts: {counter.shirt_count}")
     
-    except KeyboardInterrupt:
-        print("\n⚠ Processing interrupted by user")
+    except Exception as e:
+        cap.release()
+        out.release()
+        return None, f"❌ Error during processing: {str(e)}"
     
     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)
+    
+    progress(1.0, desc="Complete!")
+    
+    stats = counter.get_stats()
+    
+    result_text = f"""
+✅ **Processing Complete!**
+
+📊 **Results:**
+- Total Frames Processed: {frame_count:,}
+- **Shirts Counted: {stats['total_shirts']}**
+- Final State: {stats['current_state']}
 
+📝 **Event Log (Shirt Counts):**
+"""
+    
+    for evt in stats['events']:
+        if evt['event'] == 'SHIRT_COUNTED':
+            result_text += f"\n  ✓ [{evt['timestamp']}] Shirt #{evt['count']} - {evt['details']}"
+    
+    if stats['total_shirts'] == 0:
+        result_text += "\n\n⚠️ No shirts detected. Try adjusting parameters or ensure video shows the rotating pad system."
+    
+    return output_path, result_text
+
+
+# Gradio Interface
+with gr.Blocks(title="Rotating Pad Shirt Counter", theme=gr.themes.Soft()) as demo:
+    gr.Markdown("""
+    # 👕 Rotating Pad Shirt Counter
+    
+    ### Demo Showcase - Limited Training Model
+    
+    **⚠️ Important Note:** This is a demonstration model trained on only **half of a single video** for showcase purposes. 
+    Performance may vary with different videos, lighting conditions, or camera angles.
+    
+    ### How it works:
+    1. Upload a video showing a rotating pad system with shirts
+    2. The model detects when shirts are placed on the pad
+    3. System counts shirts as they rotate through the Region of Interest (ROI)
+    
+    ### Best Results:
+    - Similar camera angle and lighting to training data
+    - Clear view of the rotating pad
+    - Videos from the same or similar production line
+    
+    ---
+    """)
+    
+    with gr.Row():
+        with gr.Column():
+            video_input = gr.Video(label="Upload Video", height=400)
+            
+            with gr.Accordion("⚙️ Advanced Settings (Optional)", open=False):
+                roi_radius = gr.Slider(
+                    minimum=100, maximum=300, value=180, step=10,
+                    label="ROI Radius (pixels)",
+                    info="Detection area size around center"
+                )
+                min_confidence = gr.Slider(
+                    minimum=0.5, maximum=0.99, value=0.98, step=0.01,
+                    label="Minimum Confidence",
+                    info="Higher = more strict detection"
+                )
+                stability_frames = gr.Slider(
+                    minimum=3, maximum=30, value=15, step=1,
+                    label="Stability Frames",
+                    info="Frames needed to confirm state change"
+                )
+            
+            process_btn = gr.Button("🚀 Process Video", variant="primary", size="lg")
+        
+        with gr.Column():
+            video_output = gr.Video(label="Processed Output", height=400)
+            result_text = gr.Textbox(
+                label="Results & Statistics",
+                lines=10,
+                max_lines=15
+            )
+    
+    gr.Markdown("""
+    ---
+    ### 📌 Model Information:
+    - **Classes Detected:** `empty_pad`, `occupied_pad` (shirt on pad)
+    - **Training Data:** Half portion of single production video
+    - **Purpose:** Demonstration and proof-of-concept
+    - **Limitations:** May not generalize well to different environments
+    
+    ### 💡 Tips:
+    - Start with default settings
+    - If no shirts detected, try lowering confidence threshold
+    - If too many false counts, increase stability frames
+    - ROI radius should cover the area where pad appears
+    """)
+    
+    process_btn.click(
+        fn=process_video,
+        inputs=[video_input, roi_radius, min_confidence, stability_frames],
+        outputs=[video_output, result_text]
+    )
 
 if __name__ == "__main__":
-    process_video()
+    demo.launch()