app.py

from flask import Flask, request, jsonify
from ultralytics import YOLO
import cv2
import numpy as np
from shapely.geometry import box as shapely_box, Polygon
from shapely.ops import unary_union
import mediapipe as mp
from collections import defaultdict
from concurrent.futures import ThreadPoolExecutor, as_completed

# -------------------- CONFIG --------------------
model_paths = {
    "pores": "pores.pt",
    "pig1": "pig1.pt",       # ✅ pigmentation model
    "combine": "combine.pt", # ✅ combine model
    "wrinkle": "wrinkle.pt",
}
default_conf_threshold = 0.05
special_conf_threshold = 0.08   # ✅ for combine & pig1
pores2_conf_threshold = 0.02   # ✅ special for pores2
imgsz = 1024

# -------------------- INIT --------------------
app = Flask(__name__)
models = {name: YOLO(path) for name, path in model_paths.items()}

mp_face_mesh = mp.solutions.face_mesh
face_mesh = mp_face_mesh.FaceMesh(static_image_mode=True)

# -------------------- SKIN TYPE DETECTOR --------------------
def detect_skin_type_from_image(img):
    img = cv2.resize(img, (400, 400))
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    gray = cv2.equalizeHist(gray)

    # Oiliness
    _, highlights = cv2.threshold(gray, 220, 255, cv2.THRESH_BINARY)
    oily_score = np.sum(highlights > 0) / highlights.size

    # Dryness
    laplacian = cv2.Laplacian(gray, cv2.CV_64F)
    texture_score = laplacian.var()

    # Combination (different facial regions)
    h, w = gray.shape
    regions = {
        "forehead": gray[0:int(h*0.3), :],
        "nose": gray[int(h*0.3):int(h*0.6), int(w*0.4):int(w*0.6)],
        "cheeks": gray[int(h*0.3):int(h*0.7), int(w*0.1):int(w*0.9)]
    }
    region_oiliness = []
    for r in regions.values():
        _, r_highlights = cv2.threshold(r, 220, 255, cv2.THRESH_BINARY)
        r_score = np.sum(r_highlights > 0) / r_highlights.size
        region_oiliness.append(np.float64(r_score))
    combo_score = np.std(region_oiliness)

    # Normalize scores
    oily_norm = min(oily_score / 0.30, 1.0)
    dry_norm = min(texture_score / 6000.0, 1.0)
    combo_norm = min(combo_score * 5, 1.0)
    normal_norm = max(1.0 - (oily_norm + dry_norm + combo_norm) / 3, 0.0)

    # Percentages
    total = oily_norm + dry_norm + combo_norm + normal_norm + 1e-6
    percentages = {
        "Oily": round(100 * oily_norm / total, 2),
        "Dry": round(100 * dry_norm / total, 2),
        "Combination": round(100 * combo_norm / total, 2),
        "Normal": round(100 * normal_norm / total, 2)
    }

    # Final type
    final_type = max(percentages, key=percentages.get)
    final_value = percentages[final_type]

    return percentages, f"{final_type} ({final_value}%)"

# -------------------- HELPERS --------------------
def run_model(model_name, model, img, face_polygon, face_area):
    if model_name == "pores2":
        conf = pores2_conf_threshold
    elif model_name in ["combine", "pig1"]:
        conf = special_conf_threshold
    else:
        conf = default_conf_threshold

    results = model(img, conf=conf, imgsz=imgsz)
    boxes_xy = results[0].boxes.xyxy.cpu().numpy()
    boxes_cls = results[0].boxes.cls.cpu().numpy().astype(int)

    class_polygons = defaultdict(list)

    for i, cls_id in enumerate(boxes_cls):
        cls_name = model.names.get(cls_id, str(cls_id)).lower()
        if model_name == "combine" and cls_name == "wrinkle":
            continue

        x1, y1, x2, y2 = boxes_xy[i].astype(int)
        det_poly = shapely_box(x1, y1, x2, y2)
        if face_polygon.intersects(det_poly):
            intersection = det_poly.intersection(face_polygon)
            if intersection.area > 0:
                class_polygons[cls_id].append(intersection)

    skin_percentages = {name.lower(): 0.0 for name in model.names.values()}
    if model_name == "combine":
        if "wrinkle" in skin_percentages:
            skin_percentages.pop("wrinkle")

    for cls_id, polys in class_polygons.items():
        union_poly = unary_union(polys)
        pixels = union_poly.area
        percentage = (pixels / face_area) * 100 if face_area > 0 else 0.0
        cls_name = model.names.get(cls_id, str(cls_id)).lower()
        skin_percentages[cls_name] = round(percentage, 2)

    return skin_percentages

def normalize_and_merge(percentages):
    normalized = {}
    for cls_name, value in percentages.items():
        name = cls_name.lower()
        if name == "pore":
            name = "pores"
        elif name == "wrinkle":
            name = "wrinkles"
        elif name == "forehead":
            name = "forehead"
        elif name == "dark_circle":
            name = "dark circles"
        elif name == "acne_scar":
            name = "scar"

        if name in ["pigmentation", "melasma"]:
            normalized["pigmentation"] = normalized.get("pigmentation", 0.0) + value
        else:
            normalized[name] = value

    return normalized

# -------------------- ROUTES --------------------
@app.route("/", methods=["GET"])
def home():
    return jsonify({
        "message": "✅ Skin API is running",
        "usage": "POST one image (form-data key 'file') to /analyze"
    })

@app.route("/analyze", methods=["POST"])
def analyze():
    try:
        files = request.files.getlist("file")
        if not files or len(files) != 1:
            return jsonify({
                "success": False,
                "analysis": [],
                "error": "You must upload exactly 1 image."
            }), 400

        file = files[0]
        file_bytes = np.frombuffer(file.read(), np.uint8)
        img = cv2.imdecode(file_bytes, cv2.IMREAD_COLOR)
        if img is None:
            return jsonify({"success": False, "analysis": [], "error": "Invalid image."}), 400

        img_h, img_w = img.shape[:2]
        rgb_img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        result = face_mesh.process(rgb_img)

        if not result.multi_face_landmarks:
            return jsonify({"success": False, "analysis": [], "error": "No face detected."}), 400

        # Face polygon
        for landmarks in result.multi_face_landmarks:
            points = np.array([[int(lm.x * img_w), int(lm.y * img_h)] for lm in landmarks.landmark])
            hull = cv2.convexHull(points)
            face_area = cv2.contourArea(hull)
            face_polygon = Polygon(hull.reshape(-1, 2))
            break

        # Run YOLO models in parallel
        combined_percentages = {}
        with ThreadPoolExecutor() as executor:
            futures = {
                executor.submit(run_model, model_name, model, img, face_polygon, face_area): model_name
                for model_name, model in models.items()
            }
            for future in as_completed(futures):
                skin_percentages = future.result()
                combined_percentages.update(skin_percentages)

        final_percentages = normalize_and_merge(combined_percentages)

        # Merge wrinkles + forehead
        wrinkle_value = final_percentages.get("wrinkles", 0.0) + final_percentages.get("forehead", 0.0)
        final_percentages["wrinkles"] = round(wrinkle_value, 2)
        if "forehead" in final_percentages:
            final_percentages.pop("forehead")

        # Skin type
        skin_type_percentages, _ = detect_skin_type_from_image(img)
        final_skin_type = max(skin_type_percentages, key=skin_type_percentages.get)
        final_skin_type_str = f"{final_skin_type} ({skin_type_percentages[final_skin_type]}%)"

        # Format response
        analysis_list = [f"{cls_name.upper()}: {value}%" for cls_name, value in final_percentages.items()]
        analysis_list.append(f"SKIN TYPE: {final_skin_type_str}")

        return jsonify({"success": True, "analysis": ["\n".join(analysis_list)]})

    except Exception as e:
        return jsonify({"success": False, "analysis": [], "error": str(e)}), 500

# -------------------- RUN --------------------
if __name__ == "__main__":
    app.run(host="0.0.0.0", port=7860)