-
Notifications
You must be signed in to change notification settings - Fork 3
/
web-refer.py
117 lines (103 loc) · 4.51 KB
/
web-refer.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
112
113
114
115
116
117
from io import StringIO
from pathlib import Path
import streamlit as st
import time
from detect import detect
import os
import sys
import argparse
from PIL import Image
def get_subdirs(b='.'):
'''
Returns all sub-directories in a specific Path
'''
result = []
for d in os.listdir(b):
bd = os.path.join(b, d)
if os.path.isdir(bd):
result.append(bd)
return result
def get_detection_folder():
'''
Returns the latest folder in a runs\detect
'''
return max(get_subdirs(os.path.join('runs', 'detect')), key=os.path.getmtime)
if __name__ == '__main__':
st.title('YOLOv5 Streamlit App')
parser = argparse.ArgumentParser()
parser.add_argument('--weights', nargs='+', type=str,
default='weights/yolov5s.pt', help='model.pt path(s)')
parser.add_argument('--source', type=str,
default='data/images', help='source')
parser.add_argument('--img-size', type=int, default=640,
help='inference size (pixels)')
parser.add_argument('--conf-thres', type=float,
default=0.35, help='object confidence threshold')
parser.add_argument('--iou-thres', type=float,
default=0.45, help='IOU threshold for NMS')
parser.add_argument('--device', default='',
help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
parser.add_argument('--view-img', action='store_true',
help='display results')
parser.add_argument('--save-txt', action='store_true',
help='save results to *.txt')
parser.add_argument('--save-conf', action='store_true',
help='save confidences in --save-txt labels')
parser.add_argument('--nosave', action='store_true',
help='do not save images/videos')
parser.add_argument('--classes', nargs='+', type=int,
help='filter by class: --class 0, or --class 0 2 3')
parser.add_argument('--agnostic-nms', action='store_true',
help='class-agnostic NMS')
parser.add_argument('--augment', action='store_true',
help='augmented inference')
parser.add_argument('--update', action='store_true',
help='update all models')
parser.add_argument('--project', default='runs/detect',
help='save results to project/name')
parser.add_argument('--name', default='exp',
help='save results to project/name')
parser.add_argument('--exist-ok', action='store_true',
help='existing project/name ok, do not increment')
opt = parser.parse_args()
print(opt)
source = ("Detection-image", "Detection-video")
source_index = st.sidebar.selectbox("Select your input", range(
len(source)), format_func=lambda x: source[x])
if source_index == 0:
uploaded_file = st.sidebar.file_uploader(
"Load_image", type=['png', 'jpeg', 'jpg'])
if uploaded_file is not None:
is_valid = True
with st.spinner(text='load_image...'):
st.sidebar.image(uploaded_file)
picture = Image.open(uploaded_file)
picture = picture.save(f'data/images/{uploaded_file.name}')
opt.source = f'data/images/{uploaded_file.name}'
else:
is_valid = False
else:
uploaded_file = st.sidebar.file_uploader("load_video", type=['mp4'])
if uploaded_file is not None:
is_valid = True
with st.spinner(text='load_video...'):
st.sidebar.video(uploaded_file)
with open(os.path.join("data", "videos", uploaded_file.name), "wb") as f:
f.write(uploaded_file.getbuffer())
opt.source = f'data/videos/{uploaded_file.name}'
else:
is_valid = False
if is_valid:
print('valid')
if st.button('Detect'):
# detect(opt)
if source_index == 0:
with st.spinner(text='Preparing Images'):
for img in os.listdir(get_detection_folder()):
st.image(str(Path(f'{get_detection_folder()}') / img))
st.balloons()
else:
with st.spinner(text='Preparing Video'):
for vid in os.listdir(get_detection_folder()):
st.video(str(Path(f'{get_detection_folder()}') / vid))
st.balloons()