Skip to content

Commit 3182fbe

Browse files
xugaoxiangxugaoxiang
committed
exception
1 parent bd4c032 commit 3182fbe

File tree

1 file changed

+96
-44
lines changed

1 file changed

+96
-44
lines changed

‎main.py

Lines changed: 96 additions & 44 deletions
Original file line numberDiff line numberDiff line change
@@ -34,23 +34,41 @@ def __init__(self, parent=None):
3434
# self.out = cv2.VideoWriter('prediction.avi', cv2.VideoWriter_fourcc(*'XVID'), 20.0, (640, 480))
3535

3636
parser = argparse.ArgumentParser()
37-
parser.add_argument('--weights', nargs='+', type=str, default='weights/yolov5s.pt', help='model.pt path(s)')
38-
parser.add_argument('--source', type=str, default='data/images', help='source') # file/folder, 0 for webcam
39-
parser.add_argument('--img-size', type=int, default=640, help='inference size (pixels)')
40-
parser.add_argument('--conf-thres', type=float, default=0.25, help='object confidence threshold')
41-
parser.add_argument('--iou-thres', type=float, default=0.45, help='IOU threshold for NMS')
42-
parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
43-
parser.add_argument('--view-img', action='store_true', help='display results')
44-
parser.add_argument('--save-txt', action='store_true', help='save results to *.txt')
45-
parser.add_argument('--save-conf', action='store_true', help='save confidences in --save-txt labels')
46-
parser.add_argument('--nosave', action='store_true', help='do not save images/videos')
47-
parser.add_argument('--classes', nargs='+', type=int, help='filter by class: --class 0, or --class 0 2 3')
48-
parser.add_argument('--agnostic-nms', action='store_true', help='class-agnostic NMS')
49-
parser.add_argument('--augment', action='store_true', help='augmented inference')
50-
parser.add_argument('--update', action='store_true', help='update all models')
51-
parser.add_argument('--project', default='runs/detect', help='save results to project/name')
52-
parser.add_argument('--name', default='exp', help='save results to project/name')
53-
parser.add_argument('--exist-ok', action='store_true', help='existing project/name ok, do not increment')
37+
parser.add_argument('--weights', nargs='+', type=str,
38+
default='weights/yolov5s.pt', help='model.pt path(s)')
39+
# file/folder, 0 for webcam
40+
parser.add_argument('--source', type=str,
41+
default='data/images', help='source')
42+
parser.add_argument('--img-size', type=int,
43+
default=640, help='inference size (pixels)')
44+
parser.add_argument('--conf-thres', type=float,
45+
default=0.25, help='object confidence threshold')
46+
parser.add_argument('--iou-thres', type=float,
47+
default=0.45, help='IOU threshold for NMS')
48+
parser.add_argument('--device', default='',
49+
help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
50+
parser.add_argument(
51+
'--view-img', action='store_true', help='display results')
52+
parser.add_argument('--save-txt', action='store_true',
53+
help='save results to *.txt')
54+
parser.add_argument('--save-conf', action='store_true',
55+
help='save confidences in --save-txt labels')
56+
parser.add_argument('--nosave', action='store_true',
57+
help='do not save images/videos')
58+
parser.add_argument('--classes', nargs='+', type=int,
59+
help='filter by class: --class 0, or --class 0 2 3')
60+
parser.add_argument(
61+
'--agnostic-nms', action='store_true', help='class-agnostic NMS')
62+
parser.add_argument('--augment', action='store_true',
63+
help='augmented inference')
64+
parser.add_argument('--update', action='store_true',
65+
help='update all models')
66+
parser.add_argument('--project', default='runs/detect',
67+
help='save results to project/name')
68+
parser.add_argument('--name', default='exp',
69+
help='save results to project/name')
70+
parser.add_argument('--exist-ok', action='store_true',
71+
help='existing project/name ok, do not increment')
5472
self.opt = parser.parse_args()
5573
print(self.opt)
5674

@@ -62,15 +80,18 @@ def __init__(self, parent=None):
6280
cudnn.benchmark = True
6381

6482
# Load model
65-
self.model = attempt_load(weights, map_location=self.device) # load FP32 model
83+
self.model = attempt_load(
84+
weights, map_location=self.device) # load FP32 model
6685
stride = int(self.model.stride.max()) # model stride
6786
self.imgsz = check_img_size(imgsz, s=stride) # check img_size
6887
if self.half:
6988
self.model.half() # to FP16
7089

7190
# Get names and colors
72-
self.names = self.model.module.names if hasattr(self.model, 'module') else self.model.names
73-
self.colors = [[random.randint(0, 255) for _ in range(3)] for _ in self.names]
91+
self.names = self.model.module.names if hasattr(
92+
self.model, 'module') else self.model.names
93+
self.colors = [[random.randint(0, 255)
94+
for _ in range(3)] for _ in self.names]
7495

7596
def setupUi(self, MainWindow):
7697
MainWindow.setObjectName("MainWindow")
@@ -80,17 +101,20 @@ def setupUi(self, MainWindow):
80101
self.horizontalLayout_2 = QtWidgets.QHBoxLayout(self.centralwidget)
81102
self.horizontalLayout_2.setObjectName("horizontalLayout_2")
82103
self.horizontalLayout = QtWidgets.QHBoxLayout()
83-
self.horizontalLayout.setSizeConstraint(QtWidgets.QLayout.SetNoConstraint)
104+
self.horizontalLayout.setSizeConstraint(
105+
QtWidgets.QLayout.SetNoConstraint)
84106
self.horizontalLayout.setObjectName("horizontalLayout")
85107
self.verticalLayout = QtWidgets.QVBoxLayout()
86108
self.verticalLayout.setContentsMargins(-1, -1, 0, -1)
87109
self.verticalLayout.setSpacing(80)
88110
self.verticalLayout.setObjectName("verticalLayout")
89111
self.pushButton_img = QtWidgets.QPushButton(self.centralwidget)
90-
sizePolicy = QtWidgets.QSizePolicy(QtWidgets.QSizePolicy.Minimum, QtWidgets.QSizePolicy.MinimumExpanding)
112+
sizePolicy = QtWidgets.QSizePolicy(
113+
QtWidgets.QSizePolicy.Minimum, QtWidgets.QSizePolicy.MinimumExpanding)
91114
sizePolicy.setHorizontalStretch(0)
92115
sizePolicy.setVerticalStretch(0)
93-
sizePolicy.setHeightForWidth(self.pushButton_img.sizePolicy().hasHeightForWidth())
116+
sizePolicy.setHeightForWidth(
117+
self.pushButton_img.sizePolicy().hasHeightForWidth())
94118
self.pushButton_img.setSizePolicy(sizePolicy)
95119
self.pushButton_img.setMinimumSize(QtCore.QSize(150, 100))
96120
self.pushButton_img.setMaximumSize(QtCore.QSize(150, 100))
@@ -99,12 +123,15 @@ def setupUi(self, MainWindow):
99123
font.setPointSize(12)
100124
self.pushButton_img.setFont(font)
101125
self.pushButton_img.setObjectName("pushButton_img")
102-
self.verticalLayout.addWidget(self.pushButton_img, 0, QtCore.Qt.AlignHCenter)
126+
self.verticalLayout.addWidget(
127+
self.pushButton_img, 0, QtCore.Qt.AlignHCenter)
103128
self.pushButton_camera = QtWidgets.QPushButton(self.centralwidget)
104-
sizePolicy = QtWidgets.QSizePolicy(QtWidgets.QSizePolicy.Minimum, QtWidgets.QSizePolicy.Expanding)
129+
sizePolicy = QtWidgets.QSizePolicy(
130+
QtWidgets.QSizePolicy.Minimum, QtWidgets.QSizePolicy.Expanding)
105131
sizePolicy.setHorizontalStretch(0)
106132
sizePolicy.setVerticalStretch(0)
107-
sizePolicy.setHeightForWidth(self.pushButton_camera.sizePolicy().hasHeightForWidth())
133+
sizePolicy.setHeightForWidth(
134+
self.pushButton_camera.sizePolicy().hasHeightForWidth())
108135
self.pushButton_camera.setSizePolicy(sizePolicy)
109136
self.pushButton_camera.setMinimumSize(QtCore.QSize(150, 100))
110137
self.pushButton_camera.setMaximumSize(QtCore.QSize(150, 100))
@@ -113,12 +140,15 @@ def setupUi(self, MainWindow):
113140
font.setPointSize(12)
114141
self.pushButton_camera.setFont(font)
115142
self.pushButton_camera.setObjectName("pushButton_camera")
116-
self.verticalLayout.addWidget(self.pushButton_camera, 0, QtCore.Qt.AlignHCenter)
143+
self.verticalLayout.addWidget(
144+
self.pushButton_camera, 0, QtCore.Qt.AlignHCenter)
117145
self.pushButton_video = QtWidgets.QPushButton(self.centralwidget)
118-
sizePolicy = QtWidgets.QSizePolicy(QtWidgets.QSizePolicy.Minimum, QtWidgets.QSizePolicy.Expanding)
146+
sizePolicy = QtWidgets.QSizePolicy(
147+
QtWidgets.QSizePolicy.Minimum, QtWidgets.QSizePolicy.Expanding)
119148
sizePolicy.setHorizontalStretch(0)
120149
sizePolicy.setVerticalStretch(0)
121-
sizePolicy.setHeightForWidth(self.pushButton_video.sizePolicy().hasHeightForWidth())
150+
sizePolicy.setHeightForWidth(
151+
self.pushButton_video.sizePolicy().hasHeightForWidth())
122152
self.pushButton_video.setSizePolicy(sizePolicy)
123153
self.pushButton_video.setMinimumSize(QtCore.QSize(150, 100))
124154
self.pushButton_video.setMaximumSize(QtCore.QSize(150, 100))
@@ -127,7 +157,8 @@ def setupUi(self, MainWindow):
127157
font.setPointSize(12)
128158
self.pushButton_video.setFont(font)
129159
self.pushButton_video.setObjectName("pushButton_video")
130-
self.verticalLayout.addWidget(self.pushButton_video, 0, QtCore.Qt.AlignHCenter)
160+
self.verticalLayout.addWidget(
161+
self.pushButton_video, 0, QtCore.Qt.AlignHCenter)
131162
self.verticalLayout.setStretch(2, 1)
132163
self.horizontalLayout.addLayout(self.verticalLayout)
133164
self.label = QtWidgets.QLabel(self.centralwidget)
@@ -171,14 +202,19 @@ def button_image_open(self):
171202
print('button_image_open')
172203
name_list = []
173204

174-
img_name, _ = QtWidgets.QFileDialog.getOpenFileName(self, "打开图片", "", "*.jpg;;*.png;;All Files(*)")
205+
img_name, _ = QtWidgets.QFileDialog.getOpenFileName(
206+
self, "打开图片", "", "*.jpg;;*.png;;All Files(*)")
207+
if not img_name:
208+
return
209+
175210
img = cv2.imread(img_name)
176211
print(img_name)
177212
showimg = img
178213
with torch.no_grad():
179214
img = letterbox(img, new_shape=self.opt.img_size)[0]
180215
# Convert
181-
img = img[:, :, ::-1].transpose(2, 0, 1) # BGR to RGB, to 3x416x416
216+
# BGR to RGB, to 3x416x416
217+
img = img[:, :, ::-1].transpose(2, 0, 1)
182218
img = np.ascontiguousarray(img)
183219
img = torch.from_numpy(img).to(self.device)
184220
img = img.half() if self.half else img.float() # uint8 to fp16/32
@@ -195,26 +231,37 @@ def button_image_open(self):
195231
for i, det in enumerate(pred):
196232
if det is not None and len(det):
197233
# Rescale boxes from img_size to im0 size
198-
det[:, :4] = scale_coords(img.shape[2:], det[:, :4], showimg.shape).round()
234+
det[:, :4] = scale_coords(
235+
img.shape[2:], det[:, :4], showimg.shape).round()
199236

200237
for *xyxy, conf, cls in reversed(det):
201238
label = '%s %.2f' % (self.names[int(cls)], conf)
202239
name_list.append(self.names[int(cls)])
203-
plot_one_box(xyxy, showimg, label=label, color=self.colors[int(cls)], line_thickness=2)
240+
plot_one_box(xyxy, showimg, label=label,
241+
color=self.colors[int(cls)], line_thickness=2)
204242

205243
cv2.imwrite('prediction.jpg', showimg)
206244
self.result = cv2.cvtColor(showimg, cv2.COLOR_BGR2BGRA)
207-
self.result = cv2.resize(self.result, (640, 480), interpolation=cv2.INTER_AREA)
208-
self.QtImg = QtGui.QImage(self.result.data, self.result.shape[1], self.result.shape[0], QtGui.QImage.Format_RGB32)
245+
self.result = cv2.resize(
246+
self.result, (640, 480), interpolation=cv2.INTER_AREA)
247+
self.QtImg = QtGui.QImage(
248+
self.result.data, self.result.shape[1], self.result.shape[0], QtGui.QImage.Format_RGB32)
209249
self.label.setPixmap(QtGui.QPixmap.fromImage(self.QtImg))
210250

211251
def button_video_open(self):
212-
video_name, _ = QtWidgets.QFileDialog.getOpenFileName(self, "打开视频", "", "*.mp4;;*.avi;;All Files(*)")
252+
video_name, _ = QtWidgets.QFileDialog.getOpenFileName(
253+
self, "打开视频", "", "*.mp4;;*.avi;;All Files(*)")
254+
255+
if not video_name:
256+
return
257+
213258
flag = self.cap.open(video_name)
214259
if flag == False:
215-
QtWidgets.QMessageBox.warning(self, u"Warning", u"打开视频失败", buttons=QtWidgets.QMessageBox.Ok, defaultButton=QtWidgets.QMessageBox.Ok)
260+
QtWidgets.QMessageBox.warning(
261+
self, u"Warning", u"打开视频失败", buttons=QtWidgets.QMessageBox.Ok, defaultButton=QtWidgets.QMessageBox.Ok)
216262
else:
217-
self.out = cv2.VideoWriter('prediction.avi', cv2.VideoWriter_fourcc(*'MJPG'), 20, (int(self.cap.get(3)), int(self.cap.get(4))))
263+
self.out = cv2.VideoWriter('prediction.avi', cv2.VideoWriter_fourcc(
264+
*'MJPG'), 20, (int(self.cap.get(3)), int(self.cap.get(4))))
218265
self.timer_video.start(30)
219266
self.pushButton_video.setDisabled(True)
220267
self.pushButton_img.setDisabled(True)
@@ -225,9 +272,11 @@ def button_camera_open(self):
225272
# 默认使用第一个本地camera
226273
flag = self.cap.open(0)
227274
if flag == False:
228-
QtWidgets.QMessageBox.warning(self, u"Warning", u"打开摄像头失败", buttons=QtWidgets.QMessageBox.Ok, defaultButton=QtWidgets.QMessageBox.Ok)
275+
QtWidgets.QMessageBox.warning(
276+
self, u"Warning", u"打开摄像头失败", buttons=QtWidgets.QMessageBox.Ok, defaultButton=QtWidgets.QMessageBox.Ok)
229277
else:
230-
self.out = cv2.VideoWriter('prediction.avi', cv2.VideoWriter_fourcc(*'MJPG'), 20, (int(self.cap.get(3)), int(self.cap.get(4))))
278+
self.out = cv2.VideoWriter('prediction.avi', cv2.VideoWriter_fourcc(
279+
*'MJPG'), 20, (int(self.cap.get(3)), int(self.cap.get(4))))
231280
self.timer_video.start(30)
232281
self.pushButton_video.setDisabled(True)
233282
self.pushButton_img.setDisabled(True)
@@ -251,7 +300,8 @@ def show_video_frame(self):
251300
with torch.no_grad():
252301
img = letterbox(img, new_shape=self.opt.img_size)[0]
253302
# Convert
254-
img = img[:, :, ::-1].transpose(2, 0, 1) # BGR to RGB, to 3x416x416
303+
# BGR to RGB, to 3x416x416
304+
img = img[:, :, ::-1].transpose(2, 0, 1)
255305
img = np.ascontiguousarray(img)
256306
img = torch.from_numpy(img).to(self.device)
257307
img = img.half() if self.half else img.float() # uint8 to fp16/32
@@ -268,13 +318,15 @@ def show_video_frame(self):
268318
for i, det in enumerate(pred): # detections per image
269319
if det is not None and len(det):
270320
# Rescale boxes from img_size to im0 size
271-
det[:, :4] = scale_coords(img.shape[2:], det[:, :4], showimg.shape).round()
321+
det[:, :4] = scale_coords(
322+
img.shape[2:], det[:, :4], showimg.shape).round()
272323
# Write results
273324
for *xyxy, conf, cls in reversed(det):
274325
label = '%s %.2f' % (self.names[int(cls)], conf)
275326
name_list.append(self.names[int(cls)])
276327
print(label)
277-
plot_one_box(xyxy, showimg, label=label, color=self.colors[int(cls)], line_thickness=2)
328+
plot_one_box(
329+
xyxy, showimg, label=label, color=self.colors[int(cls)], line_thickness=2)
278330

279331
self.out.write(showimg)
280332
show = cv2.resize(showimg, (640, 480))

0 commit comments

Comments
 (0)