glenn-jocher
commented
2 months ago @rppycv hello,
Thank you for sharing your question and the detailed code snippet! It's great to see you leveraging the Ultralytics solutions for object counting with your custom model. To retrieve the in and out counters for each class separately, you can modify the ObjectCounter class to access the counts directly.
Here's an example of how you can achieve this:
- Modify the
ObjectCounterclass to store counts for each class. - Access the counts after processing each frame.
Below is an updated version of your code with these modifications:
import cv2
from ultralytics import YOLO, solutions
class CustomObjectCounter(solutions.ObjectCounter):
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
self.in_count = {name: 0 for name in self.names.values()}
self.out_count = {name: 0 for name in self.names.values()}
def start_counting(self, im0, tracks):
im0 = super().start_counting(im0, tracks)
for track in tracks:
if track['in']:
self.in_count[self.names[track['cls']]] += 1
if track['out']:
self.out_count[self.names[track['cls']]] += 1
return im0
def count_objects_in_region(video_path, output_video_path, model_path):
"""Count objects in a specific region within a video."""
model = YOLO(model_path)
cap = cv2.VideoCapture(video_path)
assert cap.isOpened(), "Error reading video file"
w, h, fps = (int(cap.get(x)) for x in (cv2.CAP_PROP_FRAME_WIDTH, cv2.CAP_PROP_FRAME_HEIGHT, cv2.CAP_PROP_FPS))
region_points = [(430, 640), (780, 550), (740, 450), (390, 530)]
fourcc = int(cap.get(cv2.CAP_PROP_FOURCC)) ### cv2.VideoWriter_fourcc(*"mp4v") BUG REPORTED
video_writer = cv2.VideoWriter(output_video_path, fourcc, fps, (w, h))
counter = CustomObjectCounter(
view_img=True, reg_pts=region_points, names=model.names, draw_tracks=True, line_thickness=2
)
while cap.isOpened():
success, im0 = cap.read()
if not success:
print("Video frame is empty or video processing has been successfully completed.")
break
tracks = model.track(im0, persist=True, show=False)
im0 = counter.start_counting(im0, tracks)
video_writer.write(im0)
cap.release()
video_writer.release()
cv2.destroyAllWindows()
# Print the counts for each class
print("In counts:", counter.in_count)
print("Out counts:", counter.out_count)
count_objects_in_region("video_INPUT.mp4", "video_OUTPUT.mp4", "best.pt")In this example, the CustomObjectCounter class extends the ObjectCounter class to maintain separate in and out counts for each class. After processing the video, the counts are printed.
Feel free to adjust the code as needed for your specific use case. If you encounter any issues, please ensure you are using the latest versions of the Ultralytics packages.
Happy coding! š
rppycv
Search before asking
Question
I have a model trained with custom dataset, named "best.pt" (train with roboflow), the counter from ultralytics work fine!. View https://docs.ultralytics.com/guides/object-counting/#real-world-applications For example, in this line of code, I can print and retrieve the class parameters
counter = solutions.ObjectCounter(view_img=True, reg_pts=region_points, names=model.names, draw_tracks=True, line_thickness=2)
print(model.names) {0: 'Higados', 1: 'vesicula'}
But I need to get from the code the in and out counters for each of the two classes separately, How can I do it?
Additional
The code is as follows
import cv2 from ultralytics import YOLO, solutions
def count_objects_in_region(video_path, output_video_path, model_path): """Count objects in a specific region within a video.""" model = YOLO(model_path) cap = cv2.VideoCapture(video_path) assert cap.isOpened(), "Error reading video file" w, h, fps = (int(cap.get(x)) for x in (cv2.CAP_PROP_FRAME_WIDTH, cv2.CAP_PROP_FRAME_HEIGHT, cv2.CAP_PROP_FPS)) region_points = [(430, 640), (780, 550), (740, 450), (390, 530)] fourcc = int(cap.get(cv2.CAP_PROP_FOURCC)) ### cv2.VideoWriter_fourcc(*"mp4v") BUG REPORTED video_writer = cv2.VideoWriter(output_video_path, fourcc, fps, (w, h)) counter = solutions.ObjectCounter( view_img=True, reg_pts=region_points, names=model.names, draw_tracks=True, line_thickness=2 )
count_objects_in_region("video_INPUT.mp4", "video_OUTPUT.mp4", "best.pt")