Update app file

app.py

@@ -1,11 +1,14 @@
 import pandas as pd
 import numpy as np
 import streamlit as st
-from PIL import Image
 import random
 import sahi.utils.file
-
+import tempfile
+import os
+from PIL import Image
+from sahi import AutoDetectionModel
 from streamlit_image_comparison import image_comparison
+from ultralyticsplus.hf_utils import download_from_hub
 
 IMAGE_TO_URL = {
     'factory_pid.png' : 'https://d1afc1j4569hs1.cloudfront.net/factory-pid.png',
@@ -24,6 +27,17 @@ st.title('P&ID Object Detection')
 st.subheader(' Identify valves and pumps with deep learning model ', divider='rainbow')
 st.caption('Developed by Deep Drawings Co.')
 
+@st.cache_resource(show_spinner=False)
+def get_model():
+    yolov8_model_path = download_from_hub('DanielCerda/pid_yolov8')
+    detection_model = AutoDetectionModel.from_pretrained(
+        model_type='yolov8',
+        model_path=yolov8_model_path,
+        confidence_threshold=0.75,
+        device="cpu",
+    )
+    return detection_model
+    
 @st.cache_data(show_spinner=False)
 def download_comparison_images():
     sahi.utils.file.download_from_url(
@@ -61,7 +75,12 @@ col1, col2, col3 = st.columns(3, gap='large')
 with col1:
     st.markdown('##### Input File')
     # set input image by upload
-    image_file = st.file_uploader("Upload your diagram", type=["pdf"])
+    uploaded_file = st.file_uploader("Upload your diagram", type="pdf")
+    if uploaded_file:
+        temp_dir = tempfile.mkdtemp()
+        path = os.path.join(temp_dir, uploaded_file.name)
+        with open(path, "wb") as f:
+            f.write(uploaded_file.getvalue())
     # set input images from examples
     def radio_func(option):
         option_to_id = {
@@ -78,7 +97,8 @@ with col1:
 with col2:
     st.markdown('##### Preview')
     # visualize input image
-    if image_file is not None:
+    if uploaded_file is not None:
+        image_file = convert_pdf_file(path=path)
         image = Image.open(image_file)
     else:
         image = sahi.utils.cv.read_image_as_pil(IMAGE_TO_URL[radio])
@@ -107,12 +127,34 @@ st.write('##')
 col1, col2, col3 = st.columns([3, 1, 3])
 with col2:
     submit = st.button("🚀 Perform Prediction")
+    
+if submit:
+    # perform prediction
+    with st.spinner(text="Downloading model weight ... "):
+        detection_model = get_model()
+        
+    image_size = 1280
+
+    with st.spinner(text="Performing prediction ... "):
+        output_1, output_2 = sahi_yolov8m_inference(
+            image,
+            detection_model,
+            image_size=image_size,
+            slice_height=slice_size,
+            slice_width=slice_size,
+            overlap_height_ratio=overlap_ratio,
+            overlap_width_ratio=overlap_ratio,
+            postprocess_match_threshold=postprocess_match_threshold
+        )
+
+    st.session_state["output_1"] = output_1
+    st.session_state["output_2"] = output_2   
 
 st.write('##')
 
 col1, col2, col3 = st.columns([1, 4, 1])
 with col2:
-    st.markdown(f"##### Object Detection Result")
+    st.markdown(f"#### Object Detection Result")
     with st.container(border = True):
         static_component = image_comparison(
             img1=st.session_state["output_1"],