Real-Time Object Counter

from ultralytics import YOLO
import cv2
from collections import defaultdict
import time

class RealTimeObjectCounter:
    """Real-time object counter using YOLOv8."""

    def __init__(self, model_path="yolov8n.pt", confidence=0.5):
        self.model = YOLO(model_path)
        self.confidence = confidence
        self.frame_counts = defaultdict(list)  # class → [count per frame]
        self.total_frames = 0

    def process_source(self, source=0, target_class=None, max_frames=None):
        """
        Count objects from camera (source=0) or video file (source="path").
        target_class: if specified, only count this class (e.g., 'person')
        """
        cap = cv2.VideoCapture(source)
        if not cap.isOpened():
            print(f"Error: Cannot open {source}")
            return

        fps = cap.get(cv2.CAP_PROP_FPS) or 30
        print(f"Processing source: {'camera' if isinstance(source, int) else source}")
        print(f"FPS: {fps:.1f} | Confidence threshold: {self.confidence}")
        if target_class:
            print(f"Counting only: '{target_class}'")

        start_time = time.time()

        while True:
            ret, frame = cap.read()
            if not ret:
                break

            if max_frames and self.total_frames >= max_frames:
                break

            # Resize for faster processing
            small = cv2.resize(frame, (640, 480))

            # Run YOLO detection
            results = self.model(small, conf=self.confidence, verbose=False)

            # Count objects in this frame
            frame_class_counts = defaultdict(int)
            for result in results:
                for box in result.boxes:
                    cls = self.model.names[int(box.cls[0])]
                    if target_class is None or cls == target_class:
                        frame_class_counts[cls] += 1

            # Record counts
            for cls, count in frame_class_counts.items():
                self.frame_counts[cls].append(count)

            # Build display frame
            annotated = results[0].plot() if not target_class else small.copy()

            # Display live stats
            y = 30
            elapsed = time.time() - start_time
            fps_live = self.total_frames / elapsed if elapsed > 0 else 0
            cv2.putText(annotated, f"FPS: {fps_live:.1f}", (10, y),
                       cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0,255,0), 2); y += 25

            for cls, count in sorted(frame_class_counts.items()):
                cv2.putText(annotated, f"{cls}: {count}", (10, y),
                           cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255,255,0), 2)
                y += 25

            cv2.imshow("Object Counter", annotated)
            self.total_frames += 1

            if cv2.waitKey(1) & 0xFF == ord("q"):
                print("Stopped by user")
                break

        cap.release()
        cv2.destroyAllWindows()
        self.print_report()

    def print_report(self):
        """Print a summary report of all detections."""
        print("\n=== DETECTION REPORT ===")
        print(f"Total frames processed: {self.total_frames}")

        if not self.frame_counts:
            print("No objects detected")
            return

        for cls, counts in sorted(self.frame_counts.items()):
            avg = sum(counts) / len(counts)
            peak = max(counts)
            present_pct = sum(1 for c in counts if c > 0) / self.total_frames * 100
            print(f"  {cls:15s}: avg={avg:.1f}/frame, peak={peak}, present in {present_pct:.1f}% of frames")

# Run the counter
counter = RealTimeObjectCounter(confidence=0.5)
# Count from webcam:
# counter.process_source(source=0, max_frames=300)
# Count from video file:
counter.process_source(source="video.mp4", max_frames=500)