Open
Description
Search before asking
- I have searched the Supervision issues and found no similar feature requests.
Question
I follow the example https://supervision.roboflow.com/latest/notebooks/count-objects-crossing-the-line/ to want to count number of car. My version is 0.24.0
Here is my code.
from pathlib import Path
import cv2
import numpy as np
from ultralytics import YOLO
import logging
from datetime import datetime
import argparse
import os
from typing import Union, Dict
import supervision as sv
class YOLOCounter:
"""Counts objects using YOLOv8 as they cross a line."""
def __init__(self, model_path: str, output_dir: str):
"""
Initialize counter with model and output paths.
Args:
model_path: Path to model file or name of pretrained model (e.g., 'yolov8n')
output_dir: Directory to save counting results
"""
try:
# self.model = YOLO(model_path)
self.model = YOLO("yolov8n.pt")
logging.info(f"Successfully loaded model: {model_path}")
except Exception as e:
raise ValueError(f"Failed to load model from {model_path}: {str(e)}")
self.output_dir = Path(output_dir)
self.output_dir.mkdir(parents=True, exist_ok=True)
self.bounding_box_annotator = sv.BoundingBoxAnnotator(thickness=1)
self.label_annotator = sv.LabelAnnotator(text_thickness=1, text_scale=0.5)
self.trace_annotator = sv.TraceAnnotator()
self.line_zone = sv.LineZone(start=sv.Point(0, 10), end=sv.Point(1296, 10))
self.line_zone_annotator = sv.LineZoneAnnotator(
thickness=2,
text_thickness=1,
text_scale=1,
text_orient_to_line=True,
text_offset=0,
display_out_count = False,
display_in_count = True,
# custom_in_text="Number of objects: ",
)
self.byte_tracker = sv.ByteTrack()
# Setup logging
self._setup_logging()
def _setup_logging(self):
"""Setup logging configuration."""
log_dir = self.output_dir / 'logs'
log_dir.mkdir(exist_ok=True)
timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
logging.basicConfig(
level=logging.INFO,
format='%(asctime)s - %(message)s',
handlers=[
logging.FileHandler(log_dir / f'counting_log_{timestamp}.txt'),
logging.StreamHandler()
]
)
def count_objects_in_video(self,
video_path: Union[str, Path],
conf: float = 0.25,
line_position: float = 0.5) -> Dict:
"""
Count objects crossing a line in video using YOLO.
Args:
video_path: Path to video file
conf: Confidence threshold
line_position: Position of the counting line (0-1, relative to frame height)
"""
video_path = Path(video_path)
if not video_path.exists():
raise FileNotFoundError(f"Video file not found: {video_path}")
video_info = sv.VideoInfo.from_video_path(video_path)
print("video_info", video_info, video_info.resolution_wh)
sv.process_video(
source_path = video_path,
target_path = self.output_dir / 'tracking_count.mp4',
callback=self.callback,
# show_progress=True,
)
def callback(self, frame: np.ndarray, index:int) -> np.ndarray:
results = self.model(frame, verbose=False, conf=0.7)[0]
detections = sv.Detections.from_ultralytics(results)
detections = self.byte_tracker.update_with_detections(detections)
labels = [
f"#{tracker_id} {self.model.names[class_id]} {confidence:0.2f}"
for confidence, class_id, tracker_id
in zip(detections.confidence, detections.class_id, detections.tracker_id)
]
print("labels", labels)
annotated_frame = frame.copy()
annotated_frame = self.trace_annotator.annotate(
scene=annotated_frame,
detections=detections)
annotated_frame = self.bounding_box_annotator.annotate(
scene=annotated_frame,
detections=detections)
annotated_frame = self.label_annotator.annotate(
scene=annotated_frame,
detections=detections,
labels=labels)
self.line_zone.trigger(detections)
print("self.line_zone", self.line_zone)
print("self.line_zone.out_count", self.line_zone.out_count)
print("self.line_zone.in_count", self.line_zone.in_count)
return self.line_zone_annotator.annotate(annotated_frame, line_counter=self.line_zone)
def is_valid_model_path(model_path):
"""Check if the provided model path is valid."""
# List of known pretrained model names
pretrained_models = ['yolov8n', 'yolov8s', 'yolov8m', 'yolov8l', 'yolov8x']
# Check if it's a pretrained model name
if any(model_path.startswith(name) for name in pretrained_models):
return True
# Check if it's a file that exists
if os.path.isfile(model_path) and model_path.endswith('.pt'):
return True
return False
def main():
parser = argparse.ArgumentParser(description='Run YOLO object counting')
parser.add_argument('--model_path', type=str,
default="yolov8n",
help='Path to model weights (.pt) or pretrained model name (e.g., yolov8n)')
parser.add_argument('--output_dir', type=str,
default="outputs/counting",
help='Directory to save counting results')
parser.add_argument('--source', type=str, required=True,
help='Path to video file')
parser.add_argument('--conf', type=float, default=0.25,
help='Confidence threshold (0-1)')
parser.add_argument('--line_position', type=float, default=0.5,
help='Position of the counting line (0-1, relative to frame height)')
args = parser.parse_args()
# Validate model path
if not is_valid_model_path(args.model_path):
if not args.model_path.endswith('.pt'):
print(f"Warning: Model path '{args.model_path}' doesn't end with .pt")
print("If this is a pretrained model name, it will be downloaded")
else:
print(f"Warning: Model file '{args.model_path}' doesn't exist")
# Initialize counter
print(f"\nInitializing counter with model: {args.model_path}")
counter = YOLOCounter(args.model_path, args.output_dir)
# Process source based on type
source_path = Path(args.source)
print(f"Processing video: {source_path}")
if source_path.is_file():
if source_path.suffix in ['.mp4', '.avi', '.mov']:
# Count objects in video
stats = counter.count_objects_in_video(
source_path,
conf=args.conf,
line_position=args.line_position
)
# print("\nCounting Statistics:")
# for key, value in stats.items():
# print(f"{key}: {value}")
else:
raise ValueError("Source must be a video file (.mp4, .avi, .mov)")
else:
raise ValueError(f"Invalid source path: {source_path}")
if __name__ == "__main__":
main()
To run the above code
python src/models/06.tracking_count.py --source outputs/videos/view35_sequence2.avi
Here is the log
Initializing counter with model: yolov8n
SupervisionWarnings: BoundingBoxAnnotator is deprecated: `BoundingBoxAnnotator` is deprecated and has been renamed to `BoxAnnotator`. `BoundingBoxAnnotator` will be removed in supervision-0.26.0.
Processing video: outputs/videos/view35_sequence2.avi
video_info VideoInfo(width=1296, height=972, fps=2, total_frames=48) (1296, 972)
labels ['#1 car 0.76']
self.line_zone.out_count 0
self.line_zone.in_count 0
labels []
self.line_zone.out_count 0
self.line_zone.in_count 0
labels []
self.line_zone.out_count 0
self.line_zone.in_count 0
labels ['#2 car 0.94']
self.line_zone.out_count 0
self.line_zone.in_count 0
labels []
self.line_zone.out_count 0
self.line_zone.in_count 0
labels ['#1 car 0.75']
self.line_zone.out_count 0
self.line_zone.in_count 0
labels ['#1 car 0.90']
self.line_zone.out_count 0
self.line_zone.in_count 0
labels ['#1 car 0.90']
self.line_zone.out_count 0
self.line_zone.in_count 0
labels ['#1 car 0.91']
self.line_zone.out_count 0
self.line_zone.in_count 0
labels ['#2 car 0.91']
self.line_zone.out_count 0
self.line_zone.in_count 0
labels ['#2 car 0.87']
self.line_zone.out_count 0
self.line_zone.in_count 0
labels []
..
As you can see, it seems that the tracking works well but the line_zone
does not update in_count
and out_count
. How can I fix this?
Additional
No response