-
Notifications
You must be signed in to change notification settings - Fork 0
/
model.py
108 lines (91 loc) · 4.29 KB
/
model.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
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
import cv2
import torch
from super_gradients.training import models
import numpy as np
import math
import yaml
from langchain.llms import LlamaCpp
from langchain import PromptTemplate, LLMChain
from huggingface_hub import hf_hub_download
def video_detection(path_x):
cap = cv2.VideoCapture(path_x)
frame_width = int(cap.get(3))
frame_height = int(cap.get(4))
device = torch.device("cuda:0") if torch.cuda.is_available() else torch.device("cpu")
model = models.get('yolo_nas_s', num_classes= 77, checkpoint_path='models/yolo-nas/ckpt_best2.pth').to(device)
count = 0
with open('data/config/data2.yaml', 'r') as file:
data = yaml.safe_load(file)
classNames = data.pop('names')
out = cv2.VideoWriter('output.avi', cv2.VideoWriter_fourcc('M', 'J', 'P', 'G'), 10, (frame_width, frame_height))
while True:
ret, frame = cap.read()
count+=1
class_final_names=[]
if ret:
result = model.predict(frame, conf=0.30)
bbox_xyxys = result.prediction.bboxes_xyxy.tolist()
confidences = result.prediction.confidence
labels = result.prediction.labels.tolist()
for (bbox_xyxy, confidence, cls) in zip(bbox_xyxys, confidences, labels):
bbox = np.array(bbox_xyxy)
x1, y1, x2, y2 = bbox[0], bbox[1], bbox[2], bbox[3]
x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2)
classname = int(cls)
class_name= classNames[classname]
class_final_names.append(class_name)
conf = math.ceil((confidence*100))/100
label = f'{class_name}{conf}'
print("Frame N", count, "", x1, y1, x2, y2)
cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 255), 3)
t_size = cv2.getTextSize(label, 0, fontScale = 1, thickness=2)[0]
c2 = x1+t_size[0], y1-t_size[1] - 3
cv2.rectangle(frame, (x1, y1), c2, [255, 144, 30], -1, cv2.LINE_AA)
cv2.putText(frame, label, (x1, y1-2), 0, 1, [255, 255, 255], thickness=1, lineType = cv2.LINE_AA)
cv2.rectangle(frame, (x1, y1), (x2, y2), (255, 0, 255), 3)
resize_frame = cv2.resize(frame, (0,0), fx=0.5, fy=0.5, interpolation=cv2.INTER_AREA)
lists = np.array(class_final_names)
unique_list = np.unique(lists)
objects_detected = ','.join(unique_list)
print(objects_detected)
yield frame, objects_detected
#out.write(frame)
#cv2.imshow("Frame", resize_frame)
#if cv2.waitKey(1) & 0xFF==ord('1'):
# break
else:
break
#out.release()
#cap.release()
#cv2.destroyAllWindows()
def image_detection(path_x):
image = cv2.imread(path_x)
device = torch.device("cuda:0") if torch.cuda.is_available() else torch.device("cpu")
model = models.get('yolo_nas_s', num_classes= 77, checkpoint_path='models/yolo-nas/ckpt_best2.pth').to(device)
with open('data/config/data2.yaml', 'r') as file:
data = yaml.safe_load(file)
classNames = data.pop('names')
result = model.predict(image, conf=0.30)
bbox_xyxys = result.prediction.bboxes_xyxy.tolist()
confidences = result.prediction.confidence
labels = result.prediction.labels.tolist()
class_final_names = []
for (bbox_xyxy, confidence, cls) in zip(bbox_xyxys, confidences, labels):
bbox = np.array(bbox_xyxy)
x1, y1, x2, y2 = bbox[0], bbox[1], bbox[2], bbox[3]
x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2)
classname = int(cls)
class_name= classNames[classname]
class_final_names.append(class_name)
conf = math.ceil((confidence*100))/100
label = f'{class_name}{conf}'
cv2.rectangle(image, (x1, y1), (x2, y2), (0, 255, 255), 3)
t_size = cv2.getTextSize(label, 0, fontScale = 1, thickness=2)[0]
c2 = x1+t_size[0], y1-t_size[1] - 3
cv2.rectangle(image, (x1, y1), c2, [255, 144, 30], -1, cv2.LINE_AA)
cv2.putText(image, label, (x1, y1-2), 0, 1, [255, 255, 255], thickness=1, lineType = cv2.LINE_AA)
lists = np.array(class_final_names)
unique_list = np.unique(lists)
objects_detected = ','.join(unique_list)
print(objects_detected)
yield image, objects_detected