asad2662/face-type-classifier

1import tensorflow as tf2from PIL import Image3import io4import numpy as np5from tensorflow.keras.applications.mobilenet_v2 import preprocess_input6 7# Define class names for face types8CLASS_NAMES = ["heart", "long", "oval", "round", "square"]9SUNGLASSES_RECOMMENDATIONS = {10    "heart": ["Aviator", "Cat-eye"],11    "long": ["Oversized", "Square/Rectangular"],12    "oval": ["Square/Rectangular", "Cat-eye"],13    "round": ["Square/Rectangular", "Cat-eye"],14    "square": ["Round/Oval", "Aviator"],15}16 17 18def preprocess(img: Image.Image) -> np.ndarray:19    """20    Preprocess the PIL image:21    - Resize to 224×22422    - Apply MobileNetV2 preprocess_input23    - Add batch dimension24    """25    img = img.resize((224, 224))26    img_array = np.array(img)27    img_array = np.expand_dims(img_array, axis=0)28    return preprocess_input(img_array)29 30 31# Load dummy model once32MODEL_PATH = "face_type_classifier.keras"33MODEL = tf.keras.models.load_model(MODEL_PATH, compile=False)34 35 36def predict(image_bytes: bytes) -> dict:37    """38    Run inference on raw image bytes.39    Returns: dict with face_type, confidence, sunglasses recommendations.40    """41    try:42        img = Image.open(io.BytesIO(image_bytes)).convert("RGB")43        arr = preprocess(img)44        outputs = MODEL.predict(arr)45        idx = int(np.argmax(outputs, axis=1)[0])46        cls = CLASS_NAMES[idx]47        confidence = float(np.max(outputs) * 100)48        return {49            "face_type": cls,50            "confidence": round(confidence, 2),51            "suggested_glasses": SUNGLASSES_RECOMMENDATIONS.get(cls, []),52        }53    except Exception as e:54        return {"error": str(e)}55