iamrobm/MOS2_Defect_Detection

1import gradio as gr2import torch3import numpy as np4from PIL import Image5import os6import cv27from huggingface_hub import hf_hub_download, HfApi8 9# Hugging Face model repo10MODEL_REPO = "iamrobm/mos2-defect-models"11 12# === CONFIG ===13MODEL_DIR = "models/"14DEFAULT_CONFIDENCE = 0.5015 16# === Model cache to avoid re-loading ===17loaded_models = {}18 19def load_model(model_rel_path):20    os.makedirs(MODEL_DIR, exist_ok=True)21    local_path = os.path.join(MODEL_DIR, model_rel_path)22 23    if not os.path.isfile(local_path):24        hf_hub_download(25            repo_id=MODEL_REPO,26            filename=model_rel_path,27            local_dir=MODEL_DIR28        )29 30    if model_rel_path not in loaded_models:31        print(f"Loading model: {model_rel_path}")32        loaded_models[model_rel_path] = torch.hub.load(33            "ultralytics/yolov5", "custom", path=local_path34        )35 36    return loaded_models[model_rel_path]37 38def detect_defects(image, model_name, conf_threshold):39    model = load_model(model_name)40    model.conf = conf_threshold41 42    results = model(np.array(image))43    detections = results.pandas().xyxy[0][["name", "confidence", "xmin", "ymin", "xmax", "ymax"]]44    result_img = np.squeeze(results.render())45    return Image.fromarray(result_img), detections46 47# === Gradio Interface ===48 49# 1) Dynamically fetch the list of .pt files from your HF Model repo:50api = HfApi()51all_files = api.list_repo_files(repo_id=MODEL_REPO)52 53# Filter only the .pt weights (ignore any folders or non-weights)54model_choices = [f for f in all_files if f.endswith(".pt")]55 56# 2) Optionally sort or reorder however you like (e.g. alphabetical):57model_choices.sort()58 59# 3) Pick the first one (or whatever default) if it exists, else None:60default_model = model_choices[0] if model_choices else None61 62iface = gr.Interface(63    fn=detect_defects,64    inputs=[65        gr.Image(type="pil", label="Upload Image"),66        gr.Dropdown(67            choices=model_choices,68            label="Choose Model",69            value=default_model70        ),71        gr.Slider(72            minimum=0.01, maximum=1.0, step=0.01,73            value=DEFAULT_CONFIDENCE, label="Confidence Threshold"74        )75    ],76    outputs=[77        gr.Image(label="Detected Image"),78        gr.Dataframe(label="Detections Table")79    ],80    title="Defect Detection in Monolayer MoS2 STEM Images",81    description=(82        "Upload an image, select a model (.pt) from the HF Hub, "83        "and adjust the confidence threshold to detect defects."84    )85)86 87iface.launch(server_name="0.0.0.0", server_port=7860)