chore(tools): show demo in live window (#1138)
Browse files- tools/demo.py +16 -12
tools/demo.py
CHANGED
@@ -211,18 +211,19 @@ def imageflow_demo(predictor, vis_folder, current_time, args):
|
|
211 |
width = cap.get(cv2.CAP_PROP_FRAME_WIDTH) # float
|
212 |
height = cap.get(cv2.CAP_PROP_FRAME_HEIGHT) # float
|
213 |
fps = cap.get(cv2.CAP_PROP_FPS)
|
214 |
-
|
215 |
-
|
216 |
-
|
217 |
-
|
218 |
-
|
219 |
-
|
220 |
-
|
221 |
-
|
222 |
-
|
223 |
-
|
224 |
-
|
225 |
-
|
|
|
226 |
while True:
|
227 |
ret_val, frame = cap.read()
|
228 |
if ret_val:
|
@@ -230,6 +231,9 @@ def imageflow_demo(predictor, vis_folder, current_time, args):
|
|
230 |
result_frame = predictor.visual(outputs[0], img_info, predictor.confthre)
|
231 |
if args.save_result:
|
232 |
vid_writer.write(result_frame)
|
|
|
|
|
|
|
233 |
ch = cv2.waitKey(1)
|
234 |
if ch == 27 or ch == ord("q") or ch == ord("Q"):
|
235 |
break
|
|
|
211 |
width = cap.get(cv2.CAP_PROP_FRAME_WIDTH) # float
|
212 |
height = cap.get(cv2.CAP_PROP_FRAME_HEIGHT) # float
|
213 |
fps = cap.get(cv2.CAP_PROP_FPS)
|
214 |
+
if args.save_result:
|
215 |
+
save_folder = os.path.join(
|
216 |
+
vis_folder, time.strftime("%Y_%m_%d_%H_%M_%S", current_time)
|
217 |
+
)
|
218 |
+
os.makedirs(save_folder, exist_ok=True)
|
219 |
+
if args.demo == "video":
|
220 |
+
save_path = os.path.join(save_folder, os.path.basename(args.path))
|
221 |
+
else:
|
222 |
+
save_path = os.path.join(save_folder, "camera.mp4")
|
223 |
+
logger.info(f"video save_path is {save_path}")
|
224 |
+
vid_writer = cv2.VideoWriter(
|
225 |
+
save_path, cv2.VideoWriter_fourcc(*"mp4v"), fps, (int(width), int(height))
|
226 |
+
)
|
227 |
while True:
|
228 |
ret_val, frame = cap.read()
|
229 |
if ret_val:
|
|
|
231 |
result_frame = predictor.visual(outputs[0], img_info, predictor.confthre)
|
232 |
if args.save_result:
|
233 |
vid_writer.write(result_frame)
|
234 |
+
else:
|
235 |
+
cv2.namedWindow("yolox", cv2.WINDOW_NORMAL)
|
236 |
+
cv2.imshow("yolox", result_frame)
|
237 |
ch = cv2.waitKey(1)
|
238 |
if ch == 27 or ch == ord("q") or ch == ord("Q"):
|
239 |
break
|