Spaces:
Sleeping
Sleeping
File size: 3,189 Bytes
26fd668 efebb94 26fd668 efebb94 26fd668 efebb94 26fd668 |
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 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 |
import streamlit as st
import torch
import cv2
import mediapipe as mp
from transformers import SwinForImageClassification, AutoFeatureExtractor
from PIL import Image
import numpy as np
# Initialize face detection
mp_face_detection = mp.solutions.face_detection.FaceDetection(
model_selection=1, min_detection_confidence=0.5)
# Initialize model and labels
@st.cache_resource
def load_model():
id2label = {0: 'Heart', 1: 'Oblong', 2: 'Oval', 3: 'Round', 4: 'Square'}
label2id = {v: k for k, v in id2label.items()}
model = SwinForImageClassification.from_pretrained(
"microsoft/swin-tiny-patch4-window7-224",
label2id=label2id,
id2label=id2label,
ignore_mismatched_sizes=True
)
model.load_state_dict(torch.load('swin.pth', map_location='cpu'))
model.eval()
return model, AutoFeatureExtractor.from_pretrained("microsoft/swin-tiny-patch4-window7-224")
model, feature_extractor = load_model()
glasses_recommendations = {
"Heart": "Rimless (tanpa bingkai bawah)",
"Oblong": "Kotak",
"Oval": "Berbagai bentuk bingkai",
"Round": "Kotak",
"Square": "Oval atau bundar"
}
def preprocess_image(image):
results = mp_face_detection.process(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
if results.detections:
detection = results.detections[0]
bbox = detection.location_data.relative_bounding_box
h, w, _ = image.shape
x1 = int(bbox.xmin * w)
y1 = int(bbox.ymin * h)
x2 = int((bbox.xmin + bbox.width) * w)
y2 = int((bbox.ymin + bbox.height) * h)
image = image[y1:y2, x1:x2]
else:
raise ValueError("No face detected")
image = cv2.resize(image, (224, 224))
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
# Use feature extractor with proper batching
inputs = feature_extractor(
images=[image], # Pass as a list of images
return_tensors="pt",
padding=True # Enable padding as a safety measure
)
return inputs['pixel_values']
def predict_face_shape(image):
image_tensor = preprocess_image(image).to(torch.device("cuda" if torch.cuda.is_available() else "cpu"))
with torch.no_grad():
outputs = model(image_tensor)
predicted_class_idx = torch.argmax(outputs.logits, dim=1).item()
return id2label[predicted_class_idx]
# Streamlit UI
st.title("Face Shape & Glasses Recommender")
st.write("Upload a face photo for shape analysis and glasses recommendations")
uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
if uploaded_file is not None:
image = Image.open(uploaded_file).convert('RGB')
img_array = np.array(image)
st.image(image, caption='Uploaded Image', use_column_width=True)
try:
with st.spinner('Analyzing...'):
prediction = predict(cv2.cvtColor(img_array, cv2.COLOR_RGB2BGR))
recommendation = glasses_recommendations[prediction]
st.success(f"Predicted Face Shape: **{prediction}**")
st.info(f"Recommended Glasses Frame: **{recommendation}**")
except Exception as e:
st.error(f"Error: {str(e)}") |