AnubhutiBhardwaj's picture
refactor: model location
f8b83a0
Raw
History Blame Contribute Delete
1.47 kB
import streamlit as st
from PIL import Image
from ultralytics import YOLO
import tempfile
import os
# Load the YOLO model
model = YOLO("yolov9e.pt")
def process_image(image_path):
# Run inference
results = model(image_path)
# Save the result image
output_path = os.path.join(tempfile.gettempdir(), "result.jpg")
results[0].save(filename=output_path)
return output_path, results[0]
def main():
st.title("Object Detection App")
uploaded_file = st.file_uploader("Choose an image", type=["jpg", "png", "jpeg"])
if uploaded_file is not None:
image = Image.open(uploaded_file)
# Save the uploaded image temporarily
temp_image_path = os.path.join(tempfile.gettempdir(), uploaded_file.name)
image.save(temp_image_path)
if st.button("Process"):
result_path, result = process_image(temp_image_path)
st.image(result_path, caption="Detected Objects", use_container_width=True)
# Display detected objects details
st.write("### Detected Objects:")
i = 1
for box in result.boxes:
x1, y1, x2, y2 = box.xyxy[0]
class_id = int(box.cls[0])
label = model.names[class_id]
st.write(
f"{i}: {label.capitalize()}, **Location:** ({x1:.2f}, {y1:.2f}, {x2:.2f}, {y2:.2f})"
)
i += 1
if __name__ == "__main__":
main()