Spaces:
Running
Running
import streamlit as st | |
import cv2 | |
import numpy as np | |
import torch | |
from torchvision import transforms | |
# Load YOLOv5 models | |
models = [] | |
models.append(torch.hub.load('ultralytics/yolov5', 'yolov5s', pretrained=True)) | |
models.append(torch.hub.load('ultralytics/yolov5', 'yolov5m', pretrained=True)) | |
models.append(torch.hub.load('ultralytics/yolov5', 'yolov5l', pretrained=True)) | |
models.append(torch.hub.load('ultralytics/yolov5', 'yolov5x', pretrained=True)) | |
# Custom CSS | |
html_style = """ | |
<style> | |
.container { | |
padding: 20px; | |
background-color: #f9f9f9; | |
border-radius: 10px; | |
box-shadow: 0 4px 8px rgba(0, 0, 0, 0.1); | |
} | |
.title { | |
color: #ff69b4; | |
font-size: 36px; | |
text-align: center; | |
margin-bottom: 30px; | |
} | |
.subheader { | |
color: #ff69b4; | |
font-size: 24px; | |
margin-top: 20px; | |
} | |
.image-container { | |
margin-top: 20px; | |
text-align: center; | |
} | |
</style> | |
""" | |
st.markdown(html_style, unsafe_allow_html=True) | |
st.markdown("<h1 class='title'>AI Skin Analyzer</h1>", unsafe_allow_html=True) | |
uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"]) | |
if uploaded_file is not None: | |
image = cv2.imdecode(np.fromstring(uploaded_file.read(), np.uint8), 1) | |
st.image(image, caption="Uploaded Image", use_column_width=True) | |
st.markdown("<h2 class='subheader'>Model Predictions:</h2>", unsafe_allow_html=True) | |
# Perform object detection for each model | |
for i, model in enumerate(models): | |
st.markdown(f"<h3 class='subheader'>Model {i+1}</h3>", unsafe_allow_html=True) | |
results = model(image) | |
results.render() | |
output_image = results.imgs[0] | |
st.image(output_image, use_column_width=True) | |