abhinav1812/face_shape_detection

1import streamlit as st
2import torch
3import torch.nn as nn
4import torchvision
5import torchvision.transforms as T
6from PIL import Image
7
8# -----------------------
9# Config
10# -----------------------
11DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
12
13CLASS_NAMES = [
14    "Heart",
15    "Oblong",
16    "Oval",
17    "Round",
18    "Square"
19]
20
21MODEL_PATH = "best_model2.pth"
22
23# -----------------------
24# Load model
25# -----------------------
26@st.cache_resource
27def load_model():
28    # ✅ use weights=None instead of pretrained=False
29    model = torchvision.models.efficientnet_b4(weights=None)
30
31    model.classifier = nn.Sequential(
32        nn.Dropout(p=0.3, inplace=True),
33        nn.Linear(model.classifier[1].in_features, len(CLASS_NAMES))
34    )
35
36    state_dict = torch.load(MODEL_PATH, map_location=DEVICE)
37    model.load_state_dict(state_dict)
38
39    model.to(DEVICE)
40    model.eval()
41    return model
42
43model = load_model()
44
45# -----------------------
46# Transforms
47# -----------------------
48transform = T.Compose([
49    T.Resize((224, 224)),
50    T.ToTensor(),
51    T.Normalize(
52        mean=[0.485, 0.456, 0.406],
53        std=[0.229, 0.224, 0.225]
54    )
55])
56
57# -----------------------
58# Prediction function
59# -----------------------
60def predict_face_shape(image: Image.Image):
61    image = transform(image).unsqueeze(0).to(DEVICE)
62
63    with torch.inference_mode():
64        outputs = model(image)
65        probs = torch.softmax(outputs, dim=1)
66        conf, pred = torch.max(probs, dim=1)
67
68    return CLASS_NAMES[pred.item()], conf.item()
69
70# -----------------------
71# Streamlit UI
72# -----------------------
73st.set_page_config(page_title="Face Shape Detector", layout="centered")
74st.title("Face Shape Detection")
75st.write("Detect your face shape using camera or image upload")
76
77option = st.radio(
78    "Choose input method:",
79    ("Use Camera", "Upload Image")
80)
81
82image = None
83
84# -------- Camera input --------
85if option == "Use Camera":
86    camera_image = st.camera_input("Take a photo")
87
88    if camera_image is not None:
89        image = Image.open(camera_image).convert("RGB")
90
91        # Fix mirrored webcam image
92        image = image.transpose(Image.FLIP_LEFT_RIGHT)
93
94# -------- Upload input --------
95else:
96    uploaded_file = st.file_uploader(
97        "Upload an image",
98        type=["jpg", "jpeg", "png"]
99    )
100
101    if uploaded_file is not None:
102        image = Image.open(uploaded_file).convert("RGB")
103
104# -------- Prediction --------
105if image is not None:
106    st.image(image, caption="Input Image", use_column_width=True)
107
108    with st.spinner("Analyzing face shape..."):
109        label, confidence = predict_face_shape(image)
110
111    st.success(f"Face Shape: {label}")
112    st.write(f"Confidence: **{confidence * 100:.2f}%**")
113