-
Notifications
You must be signed in to change notification settings - Fork 1
/
model_utils.py
111 lines (94 loc) · 4.88 KB
/
model_utils.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
109
110
111
from utils.plots import plot_one_box
from PIL import ImageColor
import subprocess
import streamlit as st
import psutil
def get_gpu_memory():
result = subprocess.check_output(
[
'nvidia-smi', '--query-gpu=memory.used',
'--format=csv,nounits,noheader'
], encoding='utf-8')
gpu_memory = [int(x) for x in result.strip().split('\n')]
return gpu_memory[0]
def color_picker_fn(classname, key):
# color_picke = st.sidebar.color_picker(f'{classname}:', '#ff0003', key=key)
color_picke = '#ff0003'
color_rgb_list = list(ImageColor.getcolor(str(color_picke), "RGB"))
color = [color_rgb_list[2], color_rgb_list[1], color_rgb_list[0]]
return color
def get_yolo(img, model_type, model, confidence, color_pick_list, class_list, draw_thick):
current_no_class = []
results = model(img)
if model_type == 'carnumber':
box = results.pandas().xyxy[0]
for i in box.index:
xmin, ymin, xmax, ymax, conf, id, class_name = int(box['xmin'][i]), int(box['ymin'][i]), int(box['xmax'][i]), \
int(box['ymax'][i]), box['confidence'][i], box['class'][i], box['name'][i]
if conf > confidence:
# plot_one_box([xmin, ymin, xmax, ymax], img, label=class_name,
# color=color_pick_list[id], line_thickness=draw_thick)
img = plot_one_box([xmin, ymin, xmax, ymax], img, label=class_name,
color=color_pick_list[id], line_thickness=draw_thick)
current_no_class.append([class_name])
if model_type == 'yolov7':
box = results.pandas().xyxy[0]
for i in box.index:
xmin, ymin, xmax, ymax, conf, id, class_name = int(box['xmin'][i]), int(box['ymin'][i]), int(box['xmax'][i]), \
int(box['ymax'][i]), box['confidence'][i], box['class'][i], box['name'][i]
if conf > confidence:
# plot_one_box([xmin, ymin, xmax, ymax], img, label=class_name,
# color=color_pick_list[id], line_thickness=draw_thick)
img = plot_one_box([xmin, ymin, xmax, ymax], img, label=class_name,
color=color_pick_list[id], line_thickness=draw_thick)
current_no_class.append([class_name])
# if model_type == 'YOLOv8':
# for result in results:
# bboxs = result.boxes.xyxy
# conf = result.boxes.conf
# cls = result.boxes.cls
# for bbox, cnf, cs in zip(bboxs, conf, cls):
# xmin = int(bbox[0])
# ymin = int(bbox[1])
# xmax = int(bbox[2])
# ymax = int(bbox[3])
# if cnf > confidence:
# plot_one_box([xmin, ymin, xmax, ymax], img, label=class_list[int(cs)],
# color=color_pick_list[int(cs)], line_thickness=draw_thick)
# current_no_class.append([class_list[int(cs)]])
return img, current_no_class
def get_system_stat(stframe1, stframe2, stframe3, fps, df_fq):
# Updating Inference results
with stframe1.container():
st.markdown("<h2>Inference Statistics</h2>", unsafe_allow_html=True)
if round(fps, 4)>1:
st.markdown(f"<h4 style='color:green;'>Frame Rate: {round(fps, 4)}</h4>", unsafe_allow_html=True)
else:
st.markdown(f"<h4 style='color:red;'>Frame Rate: {round(fps, 4)}</h4>", unsafe_allow_html=True)
with stframe2.container():
st.markdown("<h3>Detected objects in curret Frame</h3>", unsafe_allow_html=True)
st.dataframe(df_fq, use_container_width=True)
with stframe3.container():
st.markdown("<h2>System Statistics</h2>", unsafe_allow_html=True)
js1, js2, js3 = st.columns(3)
# Updating System stats
with js1:
st.markdown("<h4>Memory usage</h4>", unsafe_allow_html=True)
mem_use = psutil.virtual_memory()[2]
if mem_use > 50:
js1_text = st.markdown(f"<h5 style='color:red;'>{mem_use}%</h5>", unsafe_allow_html=True)
else:
js1_text = st.markdown(f"<h5 style='color:green;'>{mem_use}%</h5>", unsafe_allow_html=True)
with js2:
st.markdown("<h4>CPU Usage</h4>", unsafe_allow_html=True)
cpu_use = psutil.cpu_percent()
if mem_use > 50:
js2_text = st.markdown(f"<h5 style='color:red;'>{cpu_use}%</h5>", unsafe_allow_html=True)
else:
js2_text = st.markdown(f"<h5 style='color:green;'>{cpu_use}%</h5>", unsafe_allow_html=True)
with js3:
st.markdown("<h4>GPU Memory Usage</h4>", unsafe_allow_html=True)
try:
js3_text = st.markdown(f'<h5>{get_gpu_memory()} MB</h5>', unsafe_allow_html=True)
except:
js3_text = st.markdown('<h5>NA</h5>', unsafe_allow_html=True)