Joakim commited on
Commit
ec26e0a
·
1 Parent(s): 59d885b

chore(tools): show demo in live window (#1138)

Browse files
Files changed (1) hide show
  1. 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
- save_folder = os.path.join(
215
- vis_folder, time.strftime("%Y_%m_%d_%H_%M_%S", current_time)
216
- )
217
- os.makedirs(save_folder, exist_ok=True)
218
- if args.demo == "video":
219
- save_path = os.path.join(save_folder, os.path.basename(args.path))
220
- else:
221
- save_path = os.path.join(save_folder, "camera.mp4")
222
- logger.info(f"video save_path is {save_path}")
223
- vid_writer = cv2.VideoWriter(
224
- save_path, cv2.VideoWriter_fourcc(*"mp4v"), fps, (int(width), int(height))
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