diff --git a/.gitignore b/.gitignore index ffbfd6d..e5bf699 100644 --- a/.gitignore +++ b/.gitignore @@ -3,4 +3,9 @@ *.jpeg *.jpg *.mp4 -*.mpg \ No newline at end of file +*.mpg +*.csv +*.webp +*.png + +flagged*/ diff --git a/gradio_web_ui.py b/gradio_web_ui.py index b75d835..a284bcb 100644 --- a/gradio_web_ui.py +++ b/gradio_web_ui.py @@ -9,7 +9,11 @@ def inference_image(image, mlmodel_name: str, confidence: float): model = RTDETR(mlmodel_name) - results = model.predict(image, conf=confidence / 100) + results = model.predict( + source=image, + conf=confidence / 100, + verbose=False + ) annotated_frame = results[0].plot() results = results[0].cpu() @@ -29,8 +33,9 @@ def inference_image(image, mlmodel_name: str, confidence: float): return annotated_frame, csv_data -def infer_video(videos, mlmodel_name: str, confidence: float): +def infer_video(videos, mlmodel_name: str, confidence: float, progress=gr.Progress()): model = RTDETR(mlmodel_name) + model.info() output_files = [] boxes_info = [] for video in videos: @@ -40,6 +45,7 @@ def infer_video(videos, mlmodel_name: str, confidence: float): try: while cap.isOpened(): ret, frame = cap.read() + progress(cap.get(cv2.CAP_PROP_POS_FRAMES) / cap.get(cv2.CAP_PROP_FRAME_COUNT)) if not ret: break results = model.track( @@ -96,9 +102,19 @@ def infer_video(videos, mlmodel_name: str, confidence: float): [ gr.Image(type="numpy", label="Upload an Image"), gr.Dropdown( - glob.glob("./ml_model/*"), value="rtdetr-l.pt", label="ML Model", info="Will add more animals later!" + glob.glob("./ml_model/*"), + value="rtdetr-l.pt", + label="ML Model", + info="Please place the RT-DETR model in the ml_model directory under the root directory of this project! It supports extensions like .pt, .onnx, and .engine!" + ), + gr.Slider( + minimum=0, + maximum=100, + value=75, + label="Confidence", + step=5, + info="Choose between 0% and 100%" ), - gr.Slider(0, 100, value=75, label="Confidence", step=5, info="Choose between 0% and 100%"), ], [ gr.Image(type="numpy", label="result image"), @@ -109,17 +125,26 @@ def infer_video(videos, mlmodel_name: str, confidence: float): gr.Interface( infer_video, [ - gr.File(label="Upload a Video", file_count="multiple", file_types=["mp4", "mpg"]), + gr.File(label="Upload a Video", file_count="multiple", file_types=["mp4", "mpg", "MOV"]), gr.Dropdown( - glob.glob("./ml_model/*"), value="rtdetr-l.pt", label="ML Model", info="Will add more animals later!" + glob.glob("./ml_model/*"), + value="rtdetr-l.pt", + label="ML Model", + info="Please place the RT-DETR model in the ml_model directory under the root directory of this project! It supports extensions like .pt, .onnx, and .engine!" + ), + gr.Slider( + minimum=0, + maximum=100, + value=75, + label="Confidence", + step=5, + info="Choose between 0% and 100%" ), - gr.Slider(0, 100, value=75, label="Confidence", step=5, info="Choose between 0% and 100%"), ], [ gr.File(label="Annotated Video"), - # gr.Textbox(label="Bounding Boxes CSV"), ] ) if __name__ == "__main__": - main_ui.launch(server_name="0.0.0.0") + main_ui.queue().launch(server_name="0.0.0.0") diff --git a/model_download.sh b/model_download.sh index 1b77d13..e3d0436 100755 --- a/model_download.sh +++ b/model_download.sh @@ -1,3 +1,8 @@ #!/bin/bash mkdir -p ml_model + +# embryo model gdown --fuzzy -O ./ml_model/demo_model.pt 'https://docs.google.com/uc?export=download&id=1mnEUifj4NQXptB3rL48a8EarYtGrUqJx' + +# coco train model +gdown --fuzzy -O ./ml_model/coco.pt 'https://docs.google.com/uc?export=download&id=1meIWpfRB0KYe7OcGC8j63dVCZGLKGOUb' \ No newline at end of file diff --git a/run_dev_env.sh b/run_dev_env.sh index b2888a2..b8c869c 100755 --- a/run_dev_env.sh +++ b/run_dev_env.sh @@ -31,6 +31,7 @@ PLATFORM="$(uname -m)" if [ $PLATFORM = "x86_64" ]; then echo "x86" + docker pull ghcr.io/moriyalab/horus_inference_server:latest docker run -it --rm --gpus all --runtime nvidia --shm-size=32G -v $ROOT/datasets:/usr/src/datasets -v $ROOT:/home/root -w /home/root --network host ghcr.io/moriyalab/horus_inference_server:latest else echo "Not Support Platform. Only support x86."