File size: 5,270 Bytes
bd7c529
 
39debef
3170d22
 
837f8a9
 
 
 
1d178bd
4c6c8f5
bd7c529
837f8a9
0715e8c
3170d22
 
 
bd7c529
 
0715e8c
39debef
8ed86f2
3170d22
8ed86f2
3cf084f
3170d22
8ed86f2
3170d22
 
cf4de60
 
837f8a9
 
 
 
 
 
 
 
 
 
 
6c45483
 
 
 
 
 
 
 
 
 
 
 
 
bd7c529
6c45483
bd7c529
 
6c45483
bd7c529
56e0661
8c0dd07
0715e8c
76a2d16
cf4de60
ee620ac
6494a0c
4d1e216
 
76a2d16
3170d22
76a2d16
96b4bfb
4d1e216
9af4bab
 
ee620ac
4d1e216
ee620ac
 
 
 
 
 
4d1e216
 
 
 
 
 
 
 
 
ee620ac
4d1e216
 
 
9af4bab
4d1e216
 
55018fa
e422d92
55018fa
 
 
4d1e216
c72147b
9af4bab
 
9448fdb
7520ea1
 
1fc6f1f
 
82792a7
 
9af4bab
9448fdb
 
 
 
82792a7
 
9448fdb
56e0661
 
 
b574fc0
56e0661
 
837f8a9
 
 
bed94f9
837f8a9
 
 
 
 
 
 
 
 
986adf5
 
 
 
837f8a9
 
 
 
 
bd7c529
 
 
0b13976
 
 
 
 
 
 
 
 
eff795b
0b13976
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
import pandas as pd
import numpy as np
import streamlit as st
import random
import sahi.utils.file
import tempfile
import os
from PIL import Image
from sahi import AutoDetectionModel
#from utils import convert_pdf_file
from utils import sahi_yolov8m_inference
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',
    'plant_pid.png' : 'https://d1afc1j4569hs1.cloudfront.net/plant-pid.png',
    'processing_pid.png' : 'https://d1afc1j4569hs1.cloudfront.net/processing-pid.png',
    'prediction_visual.png' : 'https://d1afc1j4569hs1.cloudfront.net/prediction_visual.png'
    }

st.set_page_config(
    page_title="P&ID Object Detection",
    layout="wide",
    initial_sidebar_state="expanded"
    )

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(
        'https://d1afc1j4569hs1.cloudfront.net/plant-pid.png',
        'plant_pid.png',
    )
    sahi.utils.file.download_from_url(
        'https://d1afc1j4569hs1.cloudfront.net/prediction_visual.png',
        'prediction_visual.png',
    )

download_comparison_images()

if "output_1" not in st.session_state:
    st.session_state["output_1"] = Image.open('plant_pid.png')

if "output_2" not in st.session_state:
    st.session_state["output_2"] = Image.open('prediction_visual.png')

col1, col2, col3 = st.columns(3, gap='medium')
with col1:
    with st.expander('How to use it'):
        st.markdown(
        '''
        1) Select any example diagram  πŸ“¬
        2) Set confidence threshold πŸ“ˆ
        3) Press to perform inference   πŸš€
        4) Visualize model predictions  πŸ”Ž
        '''
        )   

st.write('##')
   
col1, col2, col3 = st.columns(3, gap='large')
with col1:
    #st.markdown('##### Input File')
    # set input image by upload
    #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 = {
            'factory_pid.png' : 'A',
            'plant_pid.png' : 'B',
            'processing_pid.png' : 'C',
        }
        return option_to_id[option]
    radio = st.radio(
        'Select from the following examples',
        options = ['factory_pid.png', 'plant_pid.png', 'processing_pid.png'],
        format_func = radio_func,
    )
with col2:
    st.markdown('##### Preview')
    # visualize input image
    #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])
    with st.container(border = True):
        st.image(image, use_column_width = True)
        
with col3:
    st.markdown('##### Set model parameters')
    postprocess_match_threshold = st.slider(
        label = 'Select confidence threshold',
        min_value = 0.0,
        max_value = 1.0,
        value = 0.75,
        step = 0.25
    )
    postprocess_match_metric = st.slider(
        label = 'Select IoU threshold',
        min_value = 0.0,
        max_value = 1.0,
        value = 0.75,
        step = 0.25
    )

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 weights ... "):
        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")
    with st.container(border = True):
        static_component = image_comparison(
        img1=st.session_state["output_1"],
        img2=st.session_state["output_2"],
        label1='Uploaded Diagram',
        label2='Model Inference',
        width=820,
        starting_position=50,
        show_labels=True,
        make_responsive=True,
        in_memory=True,
        )