How to use Ultralytics computer vision model in app

[Yaandle]

I have been locally running a streamlit app that uses Ultralytics to test different YOLO models, it works successfully when using pre-defined models but I am trying to create an option for users to upload custom models.

When I upload a custom model it successfully uploads but I cant run inference, the rest of the page goes a little shaded.

How it looks with pre-defined model vs custom model upload. Note how in the 2nd image it is still loading/running but doesn’t make any progress after a while of waiting. It also covers/removes the single-image upload.

My relevant code.

“”"

Model and mode selection

st.subheader(“Select a model and task, or upload your own model.”, divider=‘rainbow’)
option = st.selectbox(‘Select a YOLOv8 model.’, (‘Strawberry’, ‘Grapes’, ‘Tomato’, ‘Custom’))
st.write(‘You selected:’, option, “Model”)

option2 = st.selectbox(‘Select a mode.’, (‘Detection’, ‘Segmentation’))
st.write(‘You selected:’, option2, “mode”)
model_paths = {
‘Strawberry’: ‘UltralyticsModels/Strawberrymodel.pt’,
‘Apple’: ‘UltralyticsModels/AppleModel.pt’,
‘Grapes’: ‘UltralyticsModels/GrapesModel.pt’,
‘Tomato’: ‘UltralyticsModels/TomatoModel.pt’,
‘Custom’: None,
}

seg_model_paths = {
‘StrawberrySeg’: ‘UltralyticsModels/StrawberrySeg.pt’,
‘AppleSeg’: ‘UltralyticsModels/AppleSeg.pt’,
‘GrapesSeg’: ‘UltralyticsModels/GrapesSeg.pt’,
‘TomatoSeg’: ‘UltralyticsModels/TomatoSeg.pt’,
‘CustomSeg’: None,
}

try:
model_path = ‘’
if option2 == ‘Detection’:
model_path = model_paths.get(option, ‘’)
elif option2 == ‘Segmentation’:
model_path = seg_model_paths.get(option + ‘Seg’, ‘’)

if option == 'Custom':
    uploaded_file = st.file_uploader("Upload your pretrained model", type=['pt'])
    if uploaded_file:
        model_path = uploaded_file.read()
        st.success("Custom model uploaded successfully.")
        try:
            model = YOLO(model_path) 
        except Exception as e:
            st.error(f"Failed to load the custom model. Error: {e}")
else:
    if not model_path:
        raise ValueError("Model path not found.")
    model = YOLO(model_path) 

except Exception as e:
st.error(f"Failed to load a model. Error: {e}")

“”"

[Yaandle]

Solved