File size: 2,259 Bytes
0468171 8873f80 0991bda 0468171 0991bda 8873f80 0991bda 8873f80 3357e99 0991bda 8873f80 0468171 0c66360 0468171 5054a1e 0468171 3ab565c 0468171 3357e99 5054a1e 3357e99 3ab565c 3357e99 0991bda |
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 |
import streamlit as st
from ultralytics import YOLO
from PIL import Image
import numpy as np
import os
import requests
# Google Drive file ID for best.pt
FILE_ID = "1VE4I-_OMoC-wzzo-udzhJ896-WdFwJjU"
FILE_URL = f"https://drive.google.com/uc?id={FILE_ID}"
MODEL_PATH = "best.pt"
# Download the model if not already present
if not os.path.exists(MODEL_PATH):
st.info("Downloading model file from Google Drive...")
try:
response = requests.get(FILE_URL, stream=True)
response.raise_for_status()
with open(MODEL_PATH, "wb") as f:
for chunk in response.iter_content(chunk_size=8192):
if chunk:
f.write(chunk)
st.success("Model file downloaded successfully!")
except Exception as e:
st.error(f"Error downloading model: {e}")
# Load the YOLO model
try:
model = YOLO(MODEL_PATH)
except Exception as e:
st.error(f"Error loading model: {e}")
# Streamlit app
st.title("Roofy v5 basic test")
st.write("Upload an image and let the model detect objects.")
# Sliders for Confidence and Overlap thresholds
confidence_threshold = st.slider(
"Confidence Threshold",
min_value=0.0,
max_value=1.0,
value=0.25,
step=0.01,
help="Set the minimum confidence score for detections.",
)
overlap_threshold = st.slider(
"Overlap Threshold",
min_value=0.0,
max_value=1.0,
value=0.45,
step=0.01,
help="Set the maximum allowable overlap for non-max suppression.",
)
uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "png", "jpeg"])
if uploaded_file:
# Read and display the image
image = Image.open(uploaded_file)
st.image(image, caption="Uploaded Image", use_container_width=True)
# Perform prediction
with st.spinner("Processing..."):
try:
results = model.predict(
np.array(image),
conf=confidence_threshold, # Set confidence threshold
iou=overlap_threshold, # Set overlap (IoU) threshold
)
st.write("Detection Results:")
st.image(results[0].plot(), caption="Detections", use_container_width=True)
except Exception as e:
st.error(f"Error during prediction: {e}")
|