stable version

app.py

@@ -3,34 +3,21 @@ import sahi.utils.mmdet
 import sahi.model
 from PIL import Image
 import random
-from utils import imagecompare
+from utils import image_compare
 from utils import sahi_mmdet_inference
-import pathlib
-import os
 
 MMDET_YOLACT_MODEL_URL = "https://download.openmmlab.com/mmdetection/v2.0/yolact/yolact_r50_1x8_coco/yolact_r50_1x8_coco_20200908-f38d58df.pth"
 MMDET_YOLOX_MODEL_URL = "https://download.openmmlab.com/mmdetection/v2.0/yolox/yolox_tiny_8x8_300e_coco/yolox_tiny_8x8_300e_coco_20210806_234250-4ff3b67e.pth"
 MMDET_FASTERRCNN_MODEL_URL = "https://download.openmmlab.com/mmdetection/v2.0/faster_rcnn/faster_rcnn_r50_fpn_2x_coco/faster_rcnn_r50_fpn_2x_coco_bbox_mAP-0.384_20200504_210434-a5d8aa15.pth"
 
-# Images
-sahi.utils.file.download_from_url(
-    "https://user-images.githubusercontent.com/34196005/142730935-2ace3999-a47b-49bb-83e0-2bdd509f1c90.jpg",
-    "apple_tree.jpg",
-)
-sahi.utils.file.download_from_url(
-    "https://user-images.githubusercontent.com/34196005/142730936-1b397756-52e5-43be-a949-42ec0134d5d8.jpg",
-    "highway.jpg",
-)
-
-sahi.utils.file.download_from_url(
-    "https://user-images.githubusercontent.com/34196005/142742871-bf485f84-0355-43a3-be86-96b44e63c3a2.jpg",
-    "highway2.jpg",
-)
-
-sahi.utils.file.download_from_url(
-    "https://user-images.githubusercontent.com/34196005/142742872-1fefcc4d-d7e6-4c43-bbb7-6b5982f7e4ba.jpg",
-    "highway3.jpg",
-)
+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-sahi.jpg": "https://user-images.githubusercontent.com/34196005/143309867-42841f5a-9181-4d22-b570-65f90f2da231.jpg",
+}
 
 
 @st.cache(allow_output_mutation=True, show_spinner=False)
@@ -72,6 +59,34 @@ def get_mmdet_model(model_name: str):
     return detection_model
 
 
+class SpinnerTexts:
+    def __init__(self):
+        self.ind_history_list = []
+        self.text_list = [
+            "Meanwhile check out [MMDetection Colab notebook of SAHI](https://colab.research.google.com/github/obss/sahi/blob/main/demo/inference_for_mmdetection.ipynb)!",
+            "Meanwhile check out [YOLOv5 Colab notebook of SAHI](https://colab.research.google.com/github/obss/sahi/blob/main/demo/inference_for_yolov5.ipynb)!",
+            "Meanwhile check out [aerial object detection with SAHI](https://blog.ml6.eu/how-to-detect-small-objects-in-very-large-images-70234bab0f98?gi=b434299595d4)!",
+            "Meanwhile check out [COCO Utilities of SAHI](https://github.com/obss/sahi/blob/main/docs/COCO.md)!",
+            "Meanwhile check out [FiftyOne utilities of SAHI](https://github.com/obss/sahi#fiftyone-utilities)!",
+            "Meanwhile [give a Github star to SAHI](https://github.com/obss/sahi/stargazers)!",
+            "Meanwhile see [how easy is to install SAHI](https://github.com/obss/sahi#getting-started)!",
+            "Meanwhile check out [Medium blogpost of SAHI](https://medium.com/codable/sahi-a-vision-library-for-performing-sliced-inference-on-large-images-small-objects-c8b086af3b80)!",
+            "Meanwhile try out [YOLOv5 HF Spaces demo of SAHI](https://huggingface.co/spaces/fcakyon/sahi-yolov5)!",
+        ]
+
+    def _store(self, ind):
+        if len(self.ind_history_list) == 6:
+            self.ind_history_list.pop(0)
+        self.ind_history_list.append(ind)
+
+    def get(self):
+        ind = 0
+        while ind in self.ind_history_list:
+            ind = random.randint(0, len(self.text_list) - 1)
+        self._store(ind)
+        return self.text_list[ind]
+
+
 st.set_page_config(
     page_title="Small Object Detection with SAHI + YOLOX",
     page_icon="🚀",
@@ -79,6 +94,19 @@ st.set_page_config(
     initial_sidebar_state="auto",
 )
 
+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"] = sahi.utils.cv.read_image_as_pil(
+        IMAGE_TO_URL["highway2-yolox.jpg"]
+    )
+
+if "output_2" not in st.session_state:
+    st.session_state["output_2"] = sahi.utils.cv.read_image_as_pil(
+        IMAGE_TO_URL["highway2-sahi.jpg"]
+    )
+
 st.markdown(
     """
     <h2 style='text-align: center'>
@@ -93,7 +121,7 @@ st.markdown(
     <p style='text-align: center'>
     <a href='https://github.com/obss/sahi' target='_blank'>SAHI Github</a> | <a href='https://github.com/open-mmlab/mmdetection/tree/master/configs/yolox' target='_blank'>YOLOX Github</a> | <a href='https://huggingface.co/spaces/fcakyon/sahi-yolov5' target='_blank'>SAHI+YOLOv5 Demo</a>
     <br />
-    Follow me on <a href='https://twitter.com/fcakyon' target='_blank'>twitter</a>, <a href='https://www.linkedin.com/in/fcakyon/' target='_blank'>linkedin</a> and <a href='https://fcakyon.medium.com/' target='_blank'>medium</a> for more..
+    Follow me for more! <a href='https://twitter.com/fcakyon' target='_blank'>twitter</a> | <a href='https://www.linkedin.com/in/fcakyon/' target='_blank'>linkedin</a> | <a href='https://fcakyon.medium.com/' target='_blank'>medium</a>
     </p>
     """,
     unsafe_allow_html=True,
@@ -105,10 +133,12 @@ col1, col2, col3 = st.columns([6, 1, 6])
 with col1:
     st.markdown(f"##### Set input image:")
 
+    # set input image by upload
     image_file = st.file_uploader(
         "Upload an image to test:", type=["jpg", "jpeg", "png"]
     )
 
+    # set input image from exapmles
     def slider_func(option):
         option_to_id = {
             "apple_tree.jpg": str(1),
@@ -122,9 +152,16 @@ with col1:
         "Or select from example images:",
         options=["apple_tree.jpg", "highway.jpg", "highway2.jpg", "highway3.jpg"],
         format_func=slider_func,
+        value="highway2.jpg",
     )
-    image = Image.open(slider)
-    st.image(image, caption=slider, width=300)
+
+    # visualize input image
+    if image_file is not None:
+        image = Image.open(image_file)
+    else:
+        image = sahi.utils.cv.read_image_as_pil(IMAGE_TO_URL[slider])
+    st.image(image, width=300)
+
 with col3:
     st.markdown(f"##### Set SAHI parameters:")
 
@@ -148,43 +185,6 @@ col1, col2, col3 = st.columns([6, 1, 6])
 with col2:
     submit = st.button("Submit")
 
-if image_file is not None:
-    image = Image.open(image_file)
-else:
-    image = Image.open(slider)
-
-
-class SpinnerTexts:
-    def __init__(self):
-        self.ind_history_list = []
-        self.text_list = [
-            "Meanwhile check out [MMDetection Colab notebook of SAHI](https://colab.research.google.com/github/obss/sahi/blob/main/demo/inference_for_mmdetection.ipynb)!",
-            "Meanwhile check out [YOLOv5 Colab notebook of SAHI](https://colab.research.google.com/github/obss/sahi/blob/main/demo/inference_for_yolov5.ipynb)!",
-            "Meanwhile check out [aerial object detection with SAHI](https://blog.ml6.eu/how-to-detect-small-objects-in-very-large-images-70234bab0f98?gi=b434299595d4)!",
-            "Meanwhile check out [COCO Utilities of SAHI](https://github.com/obss/sahi/blob/main/docs/COCO.md)!",
-            "Meanwhile check out [FiftyOne utilities of SAHI](https://github.com/obss/sahi#fiftyone-utilities)!",
-            "Meanwhile [give a Github star to SAHI](https://github.com/obss/sahi/stargazers)!",
-            "Meanwhile see [how easy is to install SAHI](https://github.com/obss/sahi#getting-started)!",
-            "Meanwhile check out [Medium blogpost of SAHI](https://medium.com/codable/sahi-a-vision-library-for-performing-sliced-inference-on-large-images-small-objects-c8b086af3b80)!",
-            "Meanwhile try out [YOLOv5 HF Spaces demo of SAHI](https://huggingface.co/spaces/fcakyon/sahi-yolov5)!",
-        ]
-
-    def _store(self, ind):
-        if len(self.ind_history_list) == 6:
-            self.ind_history_list.pop(0)
-        self.ind_history_list.append(ind)
-
-    def get(self):
-        ind = 0
-        while ind in self.ind_history_list:
-            ind = random.randint(0, len(self.text_list) - 1)
-        self._store(ind)
-        return self.text_list[ind]
-
-
-if "last_spinner_texts" not in st.session_state:
-    st.session_state["last_spinner_texts"] = SpinnerTexts()
-
 if submit:
     # perform prediction
     with st.spinner(
@@ -215,14 +215,18 @@ if submit:
             postprocess_class_agnostic=postprocess_class_agnostic,
         )
 
-    st.markdown(f"##### YOLOX Standard vs SAHI Prediction:")
-    static_component = imagecompare(
-        output_1,
-        output_2,
-        label1="YOLOX",
-        label2="SAHI+YOLOX",
-        width=700,
-        starting_position=50,
-        show_labels=True,
-        make_responsive=True,
-    )
+    st.session_state["output_1"] = output_1
+    st.session_state["output_2"] = output_2
+
+st.markdown(f"##### YOLOX Standard vs SAHI Prediction:")
+static_component = image_compare(
+    img1=st.session_state["output_1"],
+    img2=st.session_state["output_2"],
+    label1="YOLOX",
+    label2="SAHI+YOLOX",
+    width=700,
+    starting_position=50,
+    show_labels=True,
+    make_responsive=True,
+    in_memory=True,
+)

utils.py

@@ -5,6 +5,10 @@ import sahi.utils
 from PIL import Image
 import base64
 import io
+import os
+import uuid
+
+TEMP_DIR = "temp"
 
 
 def sahi_mmdet_inference(
@@ -57,14 +61,32 @@ def sahi_mmdet_inference(
 
 def pillow_to_base64(image: Image.Image):
     in_mem_file = io.BytesIO()
-    image.save(in_mem_file, format="PNG")
+    image.save(in_mem_file, format="JPEG", subsampling=0, quality=100)
     img_bytes = in_mem_file.getvalue()  # bytes
     image_str = base64.b64encode(img_bytes).decode("utf-8")
-    base64_src = f"data:image/png;base64,{image_str}"
+    base64_src = f"data:image/jpg;base64,{image_str}"
+    return base64_src
+
+
+def local_file_to_base64(image_path: str):
+    file_ = open(image_path, "rb")
+    img_bytes = file_.read()
+    image_str = base64.b64encode(img_bytes).decode("utf-8")
+    file_.close()
+    base64_src = f"data:image/jpg;base64,{image_str}"
     return base64_src
 
 
-def imagecompare(
+def pillow_local_file_to_base64(image: Image.Image):
+    # pillow to local file
+    img_path = TEMP_DIR + "/" + str(uuid.uuid4()) + ".jpg"
+    image.save(img_path, subsampling=0, quality=100)
+    # local file base64 str
+    base64_src = local_file_to_base64(img_path)
+    return base64_src
+
+
+def image_compare(
     img1: str,
     img2: str,
     label1: str = "1",
@@ -73,6 +95,7 @@ def imagecompare(
     show_labels: bool = True,
     starting_position: int = 50,
     make_responsive: bool = True,
+    in_memory=False,
 ):
     """Create a new juxtapose component.
     Parameters
@@ -93,6 +116,8 @@ def imagecompare(
         Starting position of the slider as percent (0-100)
     make_responsive: bool or None
         Enable responsive mode
+    in_memory: bool or None
+        Handle pillow to base64 conversion in memory without saving to local
     Returns
     -------
     static_component: Boolean
@@ -103,8 +128,22 @@ def imagecompare(
     h_to_w = img_height / img_width
     height = (width * h_to_w) * 0.95
 
-    img1 = pillow_to_base64(sahi.utils.cv.read_image_as_pil(img1))
-    img2 = pillow_to_base64(sahi.utils.cv.read_image_as_pil(img2))
+    img1_pillow = sahi.utils.cv.read_image_as_pil(img1)
+    img2_pillow = sahi.utils.cv.read_image_as_pil(img2)
+
+    if in_memory:
+        # create base64 str from pillow images
+        img1 = pillow_to_base64(img1_pillow)
+        img2 = pillow_to_base64(img2_pillow)
+    else:
+        # clean temp dir
+        os.makedirs(TEMP_DIR, exist_ok=True)
+        for file_ in os.listdir(TEMP_DIR):
+            if file_.endswith(".jpg"):
+                os.remove(TEMP_DIR + "/" + file_)
+        # create base64 str from pillow images
+        img1 = pillow_local_file_to_base64(img1_pillow)
+        img2 = pillow_local_file_to_base64(img2_pillow)
 
     # load css + js
     cdn_path = "https://cdn.knightlab.com/libs/juxtapose/latest"