Spaces:
Paused
Paused
| 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 -------------------- | |
| def home(): | |
| return jsonify({ | |
| "message": "✅ Skin API is running", | |
| "usage": "POST one image (form-data key 'file') to /analyze" | |
| }) | |
| 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) | |