Merge pull request #5 from moriyalab/video_creation_opt

add video creation option

gradio_web_ui.py

@@ -5,6 +5,25 @@
 import tempfile
 import io
 import glob
+import shutil
+import os
+
+mode_list = ["Create Fully Annotated Video", "Create Time-Lapse Video", "Do Not Create Video"]
+
+
+def upload_mlmodel(filepaths):
+ dest_dir = './ml_model'
+
+ if not os.path.exists(dest_dir):
+ os.makedirs(dest_dir)
+
+ for filepath in filepaths:
+ if os.path.isfile(filepath):
+ shutil.copy(filepath, dest_dir)
+ else:
+ return f'{filepath} is not found'
+
+ return "Upload complete. Please restart gradio_web_ui.py"
 
 
 def inference_image(image, mlmodel_name: str, confidence: float):
@@ -33,14 +52,22 @@ def inference_image(image, mlmodel_name: str, confidence: float):
  return annotated_frame, csv_data
 
 
-def infer_video(videos, mlmodel_name: str, confidence: float, progress=gr.Progress()):
+def infer_video(videos, mlmodel_name: str, confidence: float, mode: float, progress=gr.Progress()):
+ global mode_list
  model = RTDETR(mlmodel_name)
  model.info()
  output_files = []
  boxes_info = []
  for video in videos:
  cap = cv2.VideoCapture(video)
- output_frames = []
+ fps = float(cap.get(cv2.CAP_PROP_FPS))
+
+ if mode_list[0] == mode or mode_list[1] == mode:
+ height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+ width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+ fourcc = cv2.VideoWriter_fourcc(*'mp4v')
+ output_file = tempfile.NamedTemporaryFile(delete=False, suffix='.mp4').name
+ out = cv2.VideoWriter(output_file, fourcc, fps, (width, height))
 
  try:
  while cap.isOpened():
@@ -57,8 +84,13 @@ def infer_video(videos, mlmodel_name: str, confidence: float, progress=gr.Progre
  half=True
  )
 
- annotated_frame = results[0].plot()
- output_frames.append(annotated_frame)
+ if mode_list[0] == mode:
+ annotated_frame = results[0].plot()
+ out.write(annotated_frame)
+ elif mode_list[1] == mode and cap.get(cv2.CAP_PROP_POS_FRAMES) % int(fps) == 0:
+ annotated_frame = results[0].plot()
+ out.write(annotated_frame)
+
  results = results[0].cpu()
  for box_data in results.boxes:
  box = box_data.xywh[0]
@@ -71,20 +103,7 @@ def infer_video(videos, mlmodel_name: str, confidence: float, progress=gr.Progre
  finally:
  cap.release()
 
- if not output_frames:
- return None
-
- height, width, _ = output_frames[0].shape
- fourcc = cv2.VideoWriter_fourcc(*'mp4v')
- output_file = tempfile.NamedTemporaryFile(delete=False, suffix='.mp4').name
- out = cv2.VideoWriter(output_file, fourcc, 30.0, (width, height))
-
- try:
- for frame in output_frames:
- out.write(frame)
- finally:
- out.release()
-
+ out.release()
  output_files.append(output_file)
 
  df = pd.DataFrame(boxes_info, columns=["xmin", "ymin", "xmax", "ymax", "confidence", "label"])
@@ -96,6 +115,16 @@ def infer_video(videos, mlmodel_name: str, confidence: float, progress=gr.Progre
 
 
 with gr.Blocks() as main_ui:
+ with gr.Tab("Upload ML Model"):
+ gr.Interface(
+ upload_mlmodel,
+ [
+ gr.File(label="Upload a ml model", file_count="multiple", file_types=["pt", "onnx", "engine"])
+ ],
+ [
+ gr.Textbox(label="Result")
+ ]
+ )
  with gr.Tab("Image Inference"):
  gr.Interface(
  inference_image,
@@ -140,6 +169,11 @@ def infer_video(videos, mlmodel_name: str, confidence: float, progress=gr.Progre
  step=5,
  info="Choose between 0% and 100%"
  ),
+ gr.Radio(
+ mode_list,
+ label="Video Creation Options",
+ info="Choose the type of video to create: fully annotated, time-lapse, or none."
+ ),
  ],
  [
  gr.File(label="Annotated Video"),