Spaces:
Runtime error
Runtime error
File size: 2,659 Bytes
bf69e73 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 |
import os
import random
import cv2
from ultralytics import YOLO
from tracker import Tracker
id2label = {0: "person", 1: "bicycle", 2:"car", 3:"motocycle", 5:"bus", 7:"truck"}
palette = (2 ** 11 - 1, 2 ** 15 - 1, 2 ** 20 - 1)
# object_counter = {"person": 0, "car": 0, "motobike": 0, "bus": 0, "truck": 0, "other": 0}
object_ids = {"person": [], "bicycle": [], "car": [], "motocycle": [], "bus": [], "truck": [], "other": []}
video_path = os.path.join('.', 'video', 'NgaTu_01.mp4')
# video_out_path = os.path.join('.', 'out.mp4')
cap = cv2.VideoCapture(video_path)
ret, frame = cap.read()
cap_out = cv2.VideoWriter('out.avi', cv2.VideoWriter_fourcc(*'MJPG'), cap.get(cv2.CAP_PROP_FPS), (frame.shape[1], frame.shape[0]))
model = YOLO('model_data/yolov8m.pt')
tracker = Tracker()
colors = [(random.randint(0, 255), random.randint(0, 255), random.randint(0, 255)) for j in range(10)]
detection_threshold = 0.7
def compute_color_for_labels(label):
if label == 0: # person
color = (85, 45, 255)
elif label == 2: # Car
color = (222, 82, 175)
elif label == 3: # Motobike
color = (0, 204, 255)
elif label == 5: # Bus
color = (0, 149, 255)
else:
color = [int((p * (label ** 2 - label + 1)) % 255) for p in palette]
return tuple(color)
while ret:
results = model.predict(frame)
for result in results:
detections = []
for r in result.boxes.data.tolist():
x1, y1, x2, y2, score, class_id = r
x1 = int(x1)
x2 = int(x2)
y1 = int(y1)
y2 = int(y2)
class_id = int(class_id)
if score > detection_threshold:
detections.append([x1, y1, x2, y2, class_id, score])
tracker.update(frame, detections)
for track in tracker.tracks:
bbox = track.bbox
x1, y1, x2, y2 = bbox
track_id = track.track_id
class_id = track.class_id
cv2.rectangle(frame, (int(x1), int(y1)), (int(x2), int(y2)), (compute_color_for_labels(class_id)), 3)
label_name = id2label[class_id] if class_id in id2label.keys() else "other"
if track_id not in object_ids[label_name]:
object_ids[label_name].append(track_id)
cv2.putText(frame,f"{label_name}-{track_id}",
(int(x1) + 5, int(y1) - 5),
cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 1, cv2.LINE_AA )
cap_out.write(frame)
cv2.imshow('frame',frame)
cv2.waitKey(2)
ret, frame = cap.read()
cap.release()
cv2.destroyAllWindows()
print(object_ids)
|