Spaces:
Sleeping
Sleeping
| 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: | |
| """ | |
| Robust shirt counter for rotating pad system. | |
| Logic: Count when empty pad ENTERS the ROI (after shirt was removed) | |
| """ | |
| def __init__(self, | |
| model_path='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""" | |
| 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 = [] | |
| for result in results: | |
| boxes = result.boxes | |
| for box in boxes: | |
| 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] | |
| center_x = (x1 + x2) / 2 | |
| center_y = (y1 + y2) / 2 | |
| 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: | |
| 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 | |
| state_counts = {} | |
| for s in self.state_buffer: | |
| state_counts[s] = state_counts.get(s, 0) + 1 | |
| stable_state = max(state_counts, key=state_counts.get) | |
| 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""" | |
| 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""" | |
| has_empty, has_occupied, detections = self.detect_in_roi(frame) | |
| instant_state = self.determine_state(has_empty, has_occupied) | |
| stable_state = self.update_state_buffer(instant_state) | |
| if self.current_state == "PAD_AWAY": | |
| self.pad_away_frames += 1 | |
| else: | |
| self.pad_away_frames = 0 | |
| state_changed = (stable_state != self.current_state) | |
| if state_changed: | |
| self.prev_state = self.current_state | |
| self.current_state = stable_state | |
| 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}") | |
| 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() | |
| cv2.circle(vis, self.roi_center, self.roi_radius, (0, 255, 255), 3) | |
| cv2.circle(vis, self.roi_center, 5, (0, 255, 255), -1) | |
| for det in detections: | |
| x1, y1, x2, y2 = map(int, det['bbox']) | |
| conf = det['confidence'] | |
| cls = det['class'] | |
| in_roi = det['in_roi'] | |
| 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 = f"{cls} {conf:.2f}" | |
| if in_roi: | |
| label += " [ROI]" | |
| cv2.putText(vis, label, (x1, y1-10), | |
| cv2.FONT_HERSHEY_SIMPLEX, 0.6, color, 2) | |
| panel_height = 180 | |
| panel = np.zeros((panel_height, vis.shape[1], 3), dtype=np.uint8) | |
| cv2.putText(panel, f"SHIRTS COUNTED: {self.shirt_count}", (20, 50), | |
| cv2.FONT_HERSHEY_SIMPLEX, 1.5, (0, 255, 0), 3) | |
| 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) | |
| cv2.putText(panel, f"Instant: {instant_state}", (20, 120), | |
| cv2.FONT_HERSHEY_SIMPLEX, 0.6, (200, 200, 200), 1) | |
| 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) | |
| 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 | |
| } | |
| 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!" | |
| progress(0, desc="Opening video...") | |
| cap = cv2.VideoCapture(video_path) | |
| if not cap.isOpened(): | |
| return None, "❌ Error: Cannot open video file" | |
| 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)) | |
| roi_center = (width // 2, height // 2) | |
| progress(0.1, desc="Loading model...") | |
| counter = RotatingPadShirtCounter( | |
| model_path='best.pt', | |
| roi_center=roi_center, | |
| roi_radius=int(roi_radius), | |
| min_conf=min_confidence, | |
| stability_frames=int(stability_frames) | |
| ) | |
| 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_path, fourcc, fps, (width, output_height)) | |
| if not out.isOpened(): | |
| cap.release() | |
| return None, "❌ Error: Cannot create output video" | |
| progress(0.2, desc="Processing video...") | |
| frame_count = 0 | |
| try: | |
| while True: | |
| ret, frame = cap.read() | |
| if not ret: | |
| break | |
| frame_count += 1 | |
| vis_frame = counter.process_frame(frame) | |
| 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) | |
| out.write(vis_frame) | |
| if frame_count % 30 == 0: | |
| progress(0.2 + (frame_count / total_frames) * 0.75, | |
| desc=f"Processing: {frame_count}/{total_frames} frames | Shirts: {counter.shirt_count}") | |
| except Exception as e: | |
| cap.release() | |
| out.release() | |
| return None, f"❌ Error during processing: {str(e)}" | |
| finally: | |
| cap.release() | |
| out.release() | |
| 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__": | |
| demo.launch() |