PeopleCounter/camera.py

"""
Camera Module for USB camera capture and frame processing
Integrates face detection and zone tracking.
"""

import cv2
import threading
import time
from face_detector import FaceDetector
from zone_tracker import ZoneTracker


class Camera:
    def __init__(self, camera_index=0, process_every_n_frames=3, 
                 face_confidence=0.5, frame_width=640, frame_height=480):
        """
        Initialize camera and processing components.
        
        Args:
            camera_index: Index of the USB camera (usually 0)
            process_every_n_frames: Process face detection every N frames for performance
            face_confidence: Confidence threshold for face detection
            frame_width: Desired frame width
            frame_height: Desired frame height
        """
        self.camera_index = camera_index
        self.process_every_n_frames = process_every_n_frames
        self.frame_width = frame_width
        self.frame_height = frame_height
        
        # Initialize camera
        self.cap = cv2.VideoCapture(camera_index)
        if not self.cap.isOpened():
            raise RuntimeError(f"Failed to open camera {camera_index}")
        
        # Set camera properties
        self.cap.set(cv2.CAP_PROP_FRAME_WIDTH, frame_width)
        self.cap.set(cv2.CAP_PROP_FRAME_HEIGHT, frame_height)
        
        # Initialize face detector and zone tracker
        self.face_detector = FaceDetector(confidence_threshold=face_confidence)
        self.zone_tracker = None  # Will be initialized after first frame
        
        # Frame processing state
        self.frame_counter = 0
        self.current_frame = None
        self.processed_frame = None
        self.current_counts = {
            'total_entered': 0,
            'total_exited': 0,
            'current_occupancy': 0
        }
        
        # Thread safety
        self.lock = threading.Lock()
        self.running = False
        self.processing_thread = None
        
        # Initialize zone tracker after getting first frame dimensions
        ret, frame = self.cap.read()
        if ret:
            h, w = frame.shape[:2]
            self.zone_tracker = ZoneTracker(w)
            self.frame_width = w
            self.frame_height = h
    
    def start(self):
        """Start the camera and processing thread."""
        if self.running:
            return
        
        self.running = True
        self.processing_thread = threading.Thread(target=self._process_loop, daemon=True)
        self.processing_thread.start()
    
    def stop(self):
        """Stop the camera and processing thread."""
        self.running = False
        if self.processing_thread:
            self.processing_thread.join(timeout=2.0)
        if self.cap:
            self.cap.release()
    
    def _process_loop(self):
        """Main processing loop running in background thread."""
        while self.running:
            ret, frame = self.cap.read()
            if not ret:
                time.sleep(0.1)
                continue
            
            self.frame_counter += 1
            
            # Store current frame
            with self.lock:
                self.current_frame = frame.copy()
            
            # Process face detection every N frames
            if self.frame_counter % self.process_every_n_frames == 0:
                processed_frame, counts = self._process_frame(frame)
                with self.lock:
                    self.processed_frame = processed_frame
                    self.current_counts = counts
    
    def _process_frame(self, frame):
        """
        Process a single frame: detect faces, track zones, update counts.
        
        Args:
            frame: Input frame from camera
            
        Returns:
            Tuple of (processed_frame, counts_dict)
        """
        # Detect faces
        faces = self.face_detector.detect_faces(frame)
        
        # Track zones and update counts
        if self.zone_tracker:
            counts = self.zone_tracker.process_faces(faces)
        else:
            counts = {
                'total_entered': 0,
                'total_exited': 0,
                'current_occupancy': 0
            }
        
        # Draw zones on frame
        if self.zone_tracker:
            processed_frame = self.zone_tracker.draw_zones(frame)
        else:
            processed_frame = frame.copy()
        
        # Draw faces on frame
        processed_frame = self.face_detector.draw_faces(processed_frame, faces)
        
        # Draw count information on frame
        text_y = 60
        cv2.putText(processed_frame, f"Entered: {counts['total_entered']}", 
                   (10, text_y), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2)
        cv2.putText(processed_frame, f"Exited: {counts['total_exited']}", 
                   (10, text_y + 30), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 255), 2)
        cv2.putText(processed_frame, f"Occupancy: {counts['current_occupancy']}", 
                   (10, text_y + 60), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 0), 2)
        
        return processed_frame, counts
    
    def get_frame(self):
        """
        Get the most recent processed frame.
        
        Returns:
            JPEG encoded frame bytes, or None if no frame available
        """
        with self.lock:
            if self.processed_frame is not None:
                ret, buffer = cv2.imencode('.jpg', self.processed_frame, 
                                          [cv2.IMWRITE_JPEG_QUALITY, 85])
                if ret:
                    return buffer.tobytes()
        return None
    
    def get_counts(self):
        """
        Get current count statistics.
        
        Returns:
            Dictionary with total_entered, total_exited, current_occupancy
        """
        with self.lock:
            return self.current_counts.copy()
    
    def reset_counts(self):
        """Reset all counters."""
        with self.lock:
            if self.zone_tracker:
                self.zone_tracker.reset_counts()
            self.current_counts = {
                'total_entered': 0,
                'total_exited': 0,
                'current_occupancy': 0
            }
    
    def __enter__(self):
        """Context manager entry."""
        self.start()
        return self
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        """Context manager exit."""
        self.stop()