Update yolo model

app.py

@@ -1,23 +1,32 @@
 import streamlit as st
-import sahi.utils.file
-import sahi.utils.mmdet
-from sahi import AutoDetectionModel
 from PIL import Image
 import random
-from utils import sahi_mmdet_inference
+from sahi.utils.yolov8 import download_yolov8m_model
+from sahi import AutoDetectionModel
+from utils import sahi_yolov8m_inference
+import sahi.utils.file
 from streamlit_image_comparison import image_comparison
 
-MMDET_YOLOX_TINY_MODEL_URL = "https://huggingface.co/fcakyon/mmdet-yolox-tiny/resolve/main/yolox_tiny_8x8_300e_coco_20211124_171234-b4047906.pth"
-MMDET_YOLOX_TINY_MODEL_PATH = "yolox.pt"
-MMDET_YOLOX_TINY_CONFIG_URL = "https://huggingface.co/fcakyon/mmdet-yolox-tiny/raw/main/yolox_tiny_8x8_300e_coco.py"
-MMDET_YOLOX_TINY_CONFIG_PATH = "config.py"
+#import sahi.utils.mmdet
+
+#MMDET_YOLOX_TINY_MODEL_URL = "https://huggingface.co/fcakyon/mmdet-yolox-tiny/resolve/main/yolox_tiny_8x8_300e_coco_20211124_171234-b4047906.pth"
+#MMDET_YOLOX_TINY_MODEL_PATH = "yolox.pt"
+#MMDET_YOLOX_TINY_CONFIG_URL = "https://huggingface.co/fcakyon/mmdet-yolox-tiny/raw/main/yolox_tiny_8x8_300e_coco.py"
+#MMDET_YOLOX_TINY_CONFIG_PATH = "config.py"
+
+#YOLOV8M_MODEL_URL = "https://github.com/ultralytics/assets/releases/download/v0.0.0/yolov8m.pt"
+#YOLOV8M_MODEL_PATH = "tests/data/models/yolov8/yolov8m.pt"
+
+
+YOLOV8M_MODEL_PATH = 'models/yolov8m.pt'
+
 
 IMAGE_TO_URL = {
     "apple_tree.jpg": "https://user-images.githubusercontent.com/34196005/142730935-2ace3999-a47b-49bb-83e0-2bdd509f1c90.jpg",
     "highway.jpg": "https://user-images.githubusercontent.com/34196005/142730936-1b397756-52e5-43be-a949-42ec0134d5d8.jpg",
     "highway2.jpg": "https://user-images.githubusercontent.com/34196005/142742871-bf485f84-0355-43a3-be86-96b44e63c3a2.jpg",
     "highway3.jpg": "https://user-images.githubusercontent.com/34196005/142742872-1fefcc4d-d7e6-4c43-bbb7-6b5982f7e4ba.jpg",
-    "highway2-yolox.jpg": "https://user-images.githubusercontent.com/34196005/143309873-c0c1f31c-c42e-4a36-834e-da0a2336bb19.jpg",
+    "highway2-yolov8m.jpg": "https://user-images.githubusercontent.com/34196005/143309873-c0c1f31c-c42e-4a36-834e-da0a2336bb19.jpg",
     "highway2-sahi.jpg": "https://user-images.githubusercontent.com/34196005/143309867-42841f5a-9181-4d22-b570-65f90f2da231.jpg",
 }
 
@@ -26,7 +35,7 @@ IMAGE_TO_URL = {
 def download_comparison_images():
     sahi.utils.file.download_from_url(
         "https://user-images.githubusercontent.com/34196005/143309873-c0c1f31c-c42e-4a36-834e-da0a2336bb19.jpg",
-        "highway2-yolox.jpg",
+        "highway2-yolov8m.jpg",
     )
     sahi.utils.file.download_from_url(
         "https://user-images.githubusercontent.com/34196005/143309867-42841f5a-9181-4d22-b570-65f90f2da231.jpg",
@@ -36,19 +45,22 @@ def download_comparison_images():
 
 @st.cache_data(show_spinner=False)
 def get_model():
-    sahi.utils.file.download_from_url(
-        MMDET_YOLOX_TINY_MODEL_URL,
-        MMDET_YOLOX_TINY_MODEL_PATH,
-    )
-    sahi.utils.file.download_from_url(
-        MMDET_YOLOX_TINY_CONFIG_URL,
-        MMDET_YOLOX_TINY_CONFIG_PATH,
-    )
+    
+    #sahi.utils.file.download_from_url(
+    #    MMDET_YOLOX_TINY_MODEL_URL,
+    #    MMDET_YOLOX_TINY_MODEL_PATH,
+    #)
+    #sahi.utils.file.download_from_url(
+    #    MMDET_YOLOX_TINY_CONFIG_URL,
+    #    MMDET_YOLOX_TINY_CONFIG_PATH,
+    #)
+    
+    download_yolov8m_model(destination_path = YOLOV8M_MODEL_PATH)
 
     detection_model = AutoDetectionModel.from_pretrained(
-        model_type='mmdet',
-        model_path=MMDET_YOLOX_TINY_MODEL_PATH,
-        config_path=MMDET_YOLOX_TINY_CONFIG_PATH,
+        model_type='yolov8',
+        model_path=YOLOV8M_MODEL_PATH,
+        #config_path=MMDET_YOLOX_TINY_CONFIG_PATH,
         confidence_threshold=0.5,
         device="cpu",
     )
@@ -96,7 +108,7 @@ if "last_spinner_texts" not in st.session_state:
     st.session_state["last_spinner_texts"] = SpinnerTexts()
 
 if "output_1" not in st.session_state:
-    st.session_state["output_1"] = Image.open("highway2-yolox.jpg")
+    st.session_state["output_1"] = Image.open("highway2-yolov8m.jpg")
 
 if "output_2" not in st.session_state:
     st.session_state["output_2"] = Image.open("highway2-sahi.jpg")
@@ -181,12 +193,12 @@ with col3:
     overlap_ratio = st.number_input(
         "overlap_ratio", min_value=0.0, max_value=0.6, value=0.2, step=0.2
     )
-    postprocess_type = st.selectbox(
-        "postprocess_type", options=["NMS", "GREEDYNMM"], index=0
-    )
-    postprocess_match_metric = st.selectbox(
-        "postprocess_match_metric", options=["IOU", "IOS"], index=0
-    )
+    #postprocess_type = st.selectbox(
+    #    "postprocess_type", options=["NMS", "GREEDYNMM"], index=0
+    #)
+    #postprocess_match_metric = st.selectbox(
+    #    "postprocess_match_metric", options=["IOU", "IOS"], index=0
+    #)
     postprocess_match_threshold = st.number_input(
         "postprocess_match_threshold", value=0.5, step=0.1
     )
@@ -209,7 +221,7 @@ if submit:
     with st.spinner(
         text="Performing prediction.. " + st.session_state["last_spinner_texts"].get()
     ):
-        output_1, output_2 = sahi_mmdet_inference(
+        output_1, output_2 = sahi_yolov8m_inference(
             image,
             detection_model,
             image_size=image_size,
@@ -226,7 +238,7 @@ if submit:
     st.session_state["output_1"] = output_1
     st.session_state["output_2"] = output_2
 
-st.markdown(f"##### YOLOX Standard vs SAHI Prediction:")
+st.markdown(f"##### YOLOv8 Standard vs SAHI Prediction:")
 static_component = image_comparison(
     img1=st.session_state["output_1"],
     img2=st.session_state["output_2"],
@@ -238,11 +250,11 @@ static_component = image_comparison(
     make_responsive=True,
     in_memory=True,
 )
-st.markdown(
-    """
-    <p style='text-align: center'>
-    prepared with <a href='https://github.com/fcakyon/streamlit-image-comparison' target='_blank'>streamlit-image-comparison</a>
-    </p>
-    """,
-    unsafe_allow_html=True,
-)
+# st.markdown(
+#     """
+#     <p style='text-align: center'>
+#     prepared with <a href='https://github.com/fcakyon/streamlit-image-comparison' target='_blank'>streamlit-image-comparison</a>
+#     </p>
+#     """,
+#     unsafe_allow_html=True,
+# )

requirements.txt

@@ -2,8 +2,9 @@
 -f https://download.openmmlab.com/mmcv/dist/cpu/torch1.10.0/index.html
 torch==1.12.1+cpu
 torchvision==0.13.1+cpu
-sahi==0.11.6
+sahi==0.11.14
 mmdet==2.25.2
 mmcv-full==1.6.1
 streamlit-image-comparison==0.0.4
-streamlit==1.22.0
+streamlit==1.22.0
+ultralyticsplus==0.0.14

utils.py

@@ -6,7 +6,7 @@ from PIL import Image
 TEMP_DIR = "temp"
 
 
-def sahi_mmdet_inference(
+def sahi_yolov8m_inference(
     image,
     detection_model,
     slice_height=512,
@@ -14,10 +14,10 @@ def sahi_mmdet_inference(
     overlap_height_ratio=0.2,
     overlap_width_ratio=0.2,
     image_size=640,
-    postprocess_type="GREEDYNMM",
-    postprocess_match_metric="IOS",
+    #postprocess_type="GREEDYNMM",
+    #postprocess_match_metric="IOS",
     postprocess_match_threshold=0.5,
-    postprocess_class_agnostic=False,
+    #postprocess_class_agnostic=False,
 ):
 
     # standard inference
@@ -39,10 +39,10 @@ def sahi_mmdet_inference(
         slice_width=slice_width,
         overlap_height_ratio=overlap_height_ratio,
         overlap_width_ratio=overlap_width_ratio,
-        postprocess_type=postprocess_type,
-        postprocess_match_metric=postprocess_match_metric,
+        #postprocess_type=postprocess_type,
+        #postprocess_match_metric=postprocess_match_metric,
         postprocess_match_threshold=postprocess_match_threshold,
-        postprocess_class_agnostic=postprocess_class_agnostic,
+        #postprocess_class_agnostic=postprocess_class_agnostic,
     )
     visual_result_2 = sahi.utils.cv.visualize_object_predictions(
         image=numpy.array(image),