test-streamlit / app.py
0x70DA's picture
Update app.py
53bdaaa
import cv2
import dlib
import streamlit as st
from PIL import Image
from transformers import pipeline
import numpy as np
# Load pre-trained image classification model from transformers library
model = pipeline("image-classification", model="0x70DA/down-syndrome-classifier")
# Load face detector from dlib library
detector = dlib.get_frontal_face_detector()
# Define the prediction function
def predict(image):
img = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR) # Convert PIL Image to NumPy array
faces = detector(img)
if len(faces) > 0:
face = faces[0] # Assuming there's only one face in the image
x, y, w, h = face.left(), face.top(), face.width(), face.height()
cropped_face = img[y: y + h, x: x + w]
# Convert the cropped image to a PIL image
pil_image = Image.fromarray(cv2.cvtColor(cropped_face, cv2.COLOR_BGR2RGB))
pred = model(pil_image)
return {o["label"]: o["score"] for o in pred}
return {"No Face Detected": 0.0}
# Create the Streamlit app interface
st.title("Down Syndrome Classifier")
uploaded_image = st.file_uploader("Upload an image", type=["jpg", "jpeg", "png"])
if uploaded_image is not None:
image = Image.open(uploaded_image)
st.image(image, caption="Uploaded Image", use_column_width=True)
st.write("Classifying...")
result = predict(image)
st.write("Classification Results:")
for label, score in result.items():
st.write(f"{label}: {score:.4f}")