Upload 2 files

face_shape_model_v2.h5

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:08d08a78219924ab0d9b8d475696f8bf912fd78dc4615cb924fa21c81f2db8d6
+size 25822424

main_app.py

@@ -0,0 +1,149 @@
+import streamlit as st
+import cv2
+import mediapipe as mp
+import numpy as np
+import os
+os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
+import tensorflow as tf
+from tensorflow.keras.models import load_model
+from PIL import Image
+
+# --- Setup ---
+BASE_PATH = r"C:\Users\MANII\Desktop\AI_Hairstyle_Project"
+MODEL_PATH = os.path.join(BASE_PATH, "face_shape_model_v2.h5")
+
+# MediaPipe
+BaseOptions = mp.tasks.BaseOptions
+FaceDetector = mp.tasks.vision.FaceDetector
+FaceDetectorOptions = mp.tasks.vision.FaceDetectorOptions
+VisionRunningMode = mp.tasks.vision.RunningMode
+TFLITE_PATH = os.path.join(BASE_PATH, "blaze_face_short_range.tflite")
+
+# Class Labels
+CLASS_NAMES = {0: 'Heart', 1: 'Oblong', 2: 'Oval', 3: 'Round', 4: 'Square'}
+
+# Hairstyle Recommendations
+RECOMMENDATIONS = {
+    'Heart': {
+        'styles': ['Side Part', 'Quiff', 'Fringe'],
+        'avoid': 'Volume on top',
+        'reason': 'Chin area balanced ho jata hai'
+    },
+    'Oblong': {
+        'styles': ['Buzz Cut', 'Crop Top', 'Side Swept'],
+        'avoid': 'Long straight styles',
+        'reason': 'Face width add hoti hai'
+    },
+    'Oval': {
+        'styles': ['Any Style', 'Undercut', 'Pompadour'],
+        'avoid': 'Kuch bhi avoid nahi',
+        'reason': 'Oval face sab styles suit karta hai'
+    },
+    'Round': {
+        'styles': ['Fade', 'Mohawk', 'Textured Top'],
+        'avoid': 'Bowl cut',
+        'reason': 'Face elongated dikhta hai'
+    },
+    'Square': {
+        'styles': ['Buzz Cut', 'Crew Cut', 'Short Sides'],
+        'avoid': 'Flat top',
+        'reason': 'Strong jawline complement hoti hai'
+    }
+}
+
+# Load ML Model
+@st.cache_resource
+def load_face_model():
+    return load_model(MODEL_PATH)
+
+ml_model = load_face_model()
+
+# --- UI ---
+st.set_page_config(page_title="AI Men's Hairstyle", layout="centered")
+st.title("✂️ Men's AI Virtual Hairstyle Try-On")
+st.markdown("Photo upload karo — AI face shape detect karega aur best hairstyle suggest karega")
+
+uploaded_file = st.file_uploader("Apni Photo Upload Karein", type=["jpg", "jpeg", "png"])
+
+if uploaded_file is not None:
+    file_bytes = np.asarray(bytearray(uploaded_file.read()), dtype=np.uint8)
+    img = cv2.imdecode(file_bytes, 1)
+    img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+
+    st.image(img_rgb, caption="Uploaded Photo", width=300)
+
+    with st.spinner("AI analyze kar raha hai..."):
+
+        # --- Face Shape Prediction ---
+        img_resized = cv2.resize(img_rgb, (224, 224))
+        img_array = np.expand_dims(img_resized / 255.0, axis=0)
+        predictions = ml_model.predict(img_array)
+        predicted_class = np.argmax(predictions[0])
+        confidence = predictions[0][predicted_class] * 100
+        face_shape = CLASS_NAMES[predicted_class]
+
+    # --- Results ---
+    st.success(f" Face Shape Detected: **{face_shape}** ({confidence:.1f}% confidence)")
+
+    rec = RECOMMENDATIONS[face_shape]
+
+    col1, col2 = st.columns(2)
+    with col1:
+        st.subheader(" Recommended Styles")
+        for style in rec['styles']:
+            st.write(f"• {style}")
+        st.caption(f"Why: {rec['reason']}")
+
+    with col2:
+        st.subheader(" Avoid")
+        st.write(rec['avoid'])
+
+    # --- Virtual Try-On ---
+    st.subheader("🎭 Virtual Try-On")
+    style_choice = st.selectbox("Hairstyle choose karo:", rec['styles'] + ['Buzz Cut', 'Second Style'])
+
+    hair_file = "buzz_cut.png" if "Buzz" in style_choice else "style.png"
+    hair_path = os.path.join(BASE_PATH, hair_file)
+    hair = cv2.imread(hair_path, cv2.IMREAD_UNCHANGED)
+
+    if hair is not None:
+        options = FaceDetectorOptions(
+            base_options=BaseOptions(model_asset_path=TFLITE_PATH),
+            running_mode=VisionRunningMode.IMAGE
+        )
+        with FaceDetector.create_from_options(options) as detector:
+            mp_image = mp.Image(image_format=mp.ImageFormat.SRGB, data=img_rgb)
+            result = detector.detect(mp_image)
+
+        if result.detections:
+            detection = result.detections[0]
+            bbox = detection.bounding_box
+            h, w, _ = img_rgb.shape
+
+            face_w = int(bbox.width * 1.1)
+            face_h = int(bbox.height * 0.6)
+            hair_resized = cv2.resize(hair, (face_w, face_h))
+
+            x1 = max(0, bbox.origin_x - int(face_w * 0.1))
+            y1 = max(0, bbox.origin_y - int(face_h * 0.7))
+            x2 = min(w, x1 + face_w)
+            y2 = min(h, y1 + face_h)
+
+            output = img_rgb.copy()
+            hair_crop = hair_resized[0:(y2-y1), 0:(x2-x1)]
+
+            if hair_crop.shape[2] == 4:
+                alpha = hair_crop[:,:,3] / 255.0
+                for c in range(3):
+                    output[y1:y2, x1:x2, c] = (
+                        hair_crop[:,:,c] * alpha +
+                        output[y1:y2, x1:x2, c] * (1 - alpha)
+                    )
+            else:
+                output[y1:y2, x1:x2] = hair_crop[:,:,:3]
+
+            st.image(output, caption=f"Try-On: {style_choice}", width=300)
+        else:
+            st.warning("Face detect nahi hua try-on ke liye.")
+    else:
+        st.error(f"Hair image nahi mili: {hair_path}")