IFMedTechdemo/Skin-Analysis

1import gradio as gr2import numpy as np3import cv24import os5from huggingface_hub import hf_hub_download6import torch7import segmentation_models_pytorch as smp8import importlib.util9#import onnxruntime as ort10 11REPO_ID = "IFMedTech/Skin-Analysis"12# List of Python files and corresponding class names13PY_MODULES = {14    "dark_circles.py": "DarkCircleDetector",15    "inflammation.py": "RednessDetector",16    "texture.py": "TextureDetector",17    "skin_tone.py": "SkinToneDetector",18    "oiliness.py": "OilinessDetector", 19    "wrinkle_unet.py": "WrinkleDetector",20    "age.py": "AgePredictor"21}22# def load_model(token):23#     """Download and load ONNX model"""24#     model_path = hf_hub_download(25#         repo_id=REPO_ID,26#         filename="model/wrinkle_model.onnx.data",  # Adjust path if needed27#         token=token28#     )29    30#     # Create ONNX Runtime session31#     session = ort.InferenceSession(32#         model_path,33#         providers=['CUDAExecutionProvider', 'CPUExecutionProvider'] 34#         if torch.cuda.is_available() else ['CPUExecutionProvider']35#     )36    37#     device = "cuda" if torch.cuda.is_available() else "cpu"38#     return session, device39def load_model(token):40    repo_id = "IFMedTech/Skin-Analysis"41    filename = "model/wrinkles_unet_v1.pth"42    # token = os.environ.get("HUGGINGFACE_HUB_TOKEN")  # Set this env var with your token43    44    # if not token:45    #     raise ValueError("HUGGINGFACE_HUB_TOKEN environment variable is required for private repo access.")46    47    model_path = hf_hub_download(repo_id=repo_id, filename=filename, token=token)48    49    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")50    model = smp.Unet(51        encoder_name="resnet34",52        encoder_weights=None,53        in_channels=3,54        classes=155    )56    model.load_state_dict(torch.load(model_path, map_location=device))57    model.to(device)58    model.eval()59    return model, device60 61def dynamic_import(module_path, class_name):62    spec = importlib.util.spec_from_file_location(class_name, module_path)63    module = importlib.util.module_from_spec(spec)64    spec.loader.exec_module(module)65    return getattr(module, class_name)66 67# Dynamically download and import modules68detector_classes = {}69token = os.environ.get("HUGGINGFACE_TOKEN")70if not token:71    raise ValueError("Please set the HUGGINGFACE_TOKEN environment variable in repo secrets!")72 73for py_file, class_name in PY_MODULES.items():74    py_path = hf_hub_download(75        repo_id=REPO_ID,76        filename=py_file,77        token=token78    )79    detector_classes[class_name] = dynamic_import(py_path, class_name)80 81# --- Skin analysis function using downloaded detectors ---82def analyze_skin(image: np.ndarray, analysis_type: str) -> np.ndarray:83    output = image.copy()84    if analysis_type == "Dark Circles":85        detector = detector_classes["DarkCircleDetector"](image)86        result = detector.predict_json()87        output = detector.draw_json()88    elif analysis_type == "Redness":89        detector = detector_classes["RednessDetector"](image)90        result = detector.predict_json()91        output = result.get("overlay_image")92    elif analysis_type == "Texture":93        detector = detector_classes["TextureDetector"](image)94        result = detector.predict_json()95        # print(result)96        output = result.get("overlay_image")97    elif analysis_type == "Skin Tone":98        detector = detector_classes["SkinToneDetector"](image)99        result = detector.predict_json()100        output = result.get("output_image")101    elif analysis_type == "Oiliness":102        detector = detector_classes["OilinessDetector"](image)103        result = detector.predict_json()104        if result.get("detected"):105            output = result.get("overlay_image")106        #     print(f"Oiliness scores: {result.get('scores')}")107        # else:108        #     print(f"Oiliness detection error: {result.get('error')}")109    elif analysis_type == "Wrinkles":110        model, device = load_model(token)111        detector = detector_classes["WrinkleDetector"](image, model, device)112        result = detector.predict_json()113        if result.get("detected") is not None:114            output = detector.draw_json(result)115 116    elif analysis_type == "Skin Age":117        detector = detector_classes["AgePredictor"](image)118        result = detector.predict_json()119        output = detector.draw_json(result)120        121    return output122 123 124# --- Gradio Interface code ---125app = gr.Interface(126    fn=analyze_skin,127    inputs=[128        gr.Image(type="numpy", label="Upload your face image"),129        gr.Radio(130            ["Dark Circles", "Redness", "Texture", "Skin Tone", "Oiliness", "Wrinkles", "Skin Age"],  131            label="Select Skin Analysis Type"132        ),133    ],134    outputs=gr.Image(type="numpy", label="Analyzed Image"),135    title="Skin Analysis Demo",136    description="Upload an image and choose a skin analysis parameter.",137        examples=[["example1.jpeg"], ["example2.jpeg"]],138)139 140if __name__ == "__main__":141    app.launch(server_name="0.0.0.0", server_port=7860)142