This repository has been archived by the owner on Jun 17, 2023. It is now read-only.
-
Notifications
You must be signed in to change notification settings - Fork 0
/
main.py
74 lines (55 loc) · 2.02 KB
/
main.py
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
import cv2
from ultralytics import YOLO
import supervision as sv
url = "https://atcs-dishub.bandung.go.id:1990/Telkom/index.m3u8"
vcap = cv2.VideoCapture(url)
model = YOLO("./best.pt")
LINE_START = sv.Point(0, 0)
LINE_END = sv.Point(640, 640)
line_counter = sv.LineZone(start=LINE_START, end=LINE_END)
line_annotator = sv.LineZoneAnnotator(thickness=2, text_thickness=1, text_scale=0.5)
box_annotator = sv.BoxAnnotator(
thickness=2,
text_thickness=1,
text_scale=0.5
)
while True:
ret, frame = vcap.read()
if frame is not None:
results = model.track(frame, stream=True, tracker="bytetrack.yaml", persist=True)
for result in results:
frame = result.orig_img
detections = sv.Detections.from_yolov8(result)
if result.boxes.id is not None:
detections.tracker_id = result.boxes.id.cpu().numpy().astype(int)
labels = []
# # i dunno what happen w/ the code below, I need some time to fix that, but its not that important
# # using empty array seems fix the issue
# labels = [
# f"{tracker_id} {model.model.names[class_id]} {confidence:0.2f}"
# for _, confidence, class_id, tracker_id
# in detections
# ]
# labels = [
# f" {class_id}"
# for class_id
# in detections
# ]
frame = box_annotator.annotate(
scene=frame,
detections=detections,
labels=labels
)
line_counter.trigger(detections=detections)
line_annotator.annotate(frame=frame, line_counter=line_counter)
print(line_counter.in_count)
print(line_counter.out_count)
cv2.imshow('frame', frame)
if cv2.waitKey(22) & 0xFF == ord('q'):
break
else:
print("Frame is None")
break
vcap.release()
cv2.destroyAllWindows()
print("Video stop")