Upload folder using huggingface_hub

README.md

@@ -1,12 +1,6 @@
 ---
-title: Derm MaskHG
-emoji: 🔥
-colorFrom: red
-colorTo: green
-sdk: gradio
-sdk_version: 5.32.0
+title: derm_maskHG
 app_file: app.py
-pinned: false
+sdk: gradio
+sdk_version: 5.27.0
 ---
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -0,0 +1,147 @@
+import os
+import torch
+import numpy as np
+import cv2 # Using OpenCV for image loading/processing
+import albumentations as A
+from albumentations.pytorch import ToTensorV2
+import gradio as gr
+
+import segmentation_models_pytorch as smp
+from train_unet import UNetLitModule # Import the Lightning Module definition
+
+# --- Configuration ---
+# Option 1: Load from the Lightning Checkpoint
+# CHECKPOINT_PATH = "checkpoints/unet-derm-epoch=XX-val_iou=Y.YYYY.ckpt" # Find the best checkpoint path from training output
+# Option 2: Load from the saved state_dict
+MODEL_STATE_DICT_PATH = "unet_derm_final_model.pth"
+IMG_SIZE = (256, 256) # MUST match training image size
+DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+# --- Load Model ---
+print(f"Loading model from: {MODEL_STATE_DICT_PATH}")
+print(f"Using device: {DEVICE}")
+
+# Instantiate the base SMP model architecture
+model = smp.Unet(
+    encoder_name="resnet34",
+    encoder_weights=None, # Don't load pretrained weights, we load our trained ones
+    in_channels=3,
+    classes=1,
+)
+
+# Load the state dict saved at the end of training
+try:
+    state_dict = torch.load(MODEL_STATE_DICT_PATH, map_location=DEVICE)
+    # If the state_dict was saved directly from the `model.model` attribute in LitModule:
+    model.load_state_dict(state_dict)
+    # If you saved the entire Lightning Module state_dict, you might need to extract the model part:
+    # state_dict = torch.load(MODEL_STATE_DICT_PATH, map_location=DEVICE)['state_dict']
+    # # Filter keys if they have a prefix like 'model.'
+    # state_dict = {k.replace('model.', ''): v for k, v in state_dict.items() if k.startswith('model.')}
+    # model.load_state_dict(state_dict)
+
+except FileNotFoundError:
+    print(f"Error: Model file not found at {MODEL_STATE_DICT_PATH}")
+    print("Please ensure the training script ran successfully and the path is correct.")
+    exit()
+except Exception as e:
+    print(f"Error loading model state_dict: {e}")
+    print("Ensure the saved state_dict matches the current model architecture.")
+    exit()
+
+
+model.to(DEVICE)
+model.eval() # Set model to evaluation mode (disables dropout, batchnorm updates)
+
+# --- Inference Transforms ---
+# Should match the validation/test transforms from training (resize, normalize)
+inference_transform = A.Compose([
+    A.Resize(height=IMG_SIZE[0], width=IMG_SIZE[1]),
+    A.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),
+    ToTensorV2(),
+])
+
+# --- Segmentation Function ---
+def segment_image(input_image_np):
+    """
+    Takes a NumPy image, performs segmentation, and returns images for Gradio.
+    """
+    # 0. Input validation
+    if input_image_np is None:
+        return None, None, None
+
+    # Ensure image is RGB (Gradio might provide BGR or grayscale)
+    if len(input_image_np.shape) == 2: # Grayscale
+        input_image_np = cv2.cvtColor(input_image_np, cv2.COLOR_GRAY2RGB)
+    elif input_image_np.shape[2] == 4: # RGBA
+        input_image_np = cv2.cvtColor(input_image_np, cv2.COLOR_RGBA2RGB)
+    # Assume BGR if 3 channels, convert to RGB for consistency with training
+    # input_image_rgb = cv2.cvtColor(input_image_np, cv2.COLOR_BGR2RGB) # PIL/Gradio usually loads RGB
+    input_image_rgb = input_image_np.copy()
+
+
+    # 1. Preprocess the image
+    transformed = inference_transform(image=input_image_rgb)
+    image_tensor = transformed['image'].unsqueeze(0).to(DEVICE) # Add batch dim and send to device
+
+    # 2. Perform inference
+    with torch.no_grad():
+        logits = model(image_tensor) # Output is [1, 1, H, W] logits
+        # Apply sigmoid to get probabilities [0, 1]
+        probabilities = torch.sigmoid(logits)
+
+    # 3. Postprocess the output mask
+    # Remove batch dimension, move to CPU, convert to NumPy
+    mask_pred_np = probabilities.squeeze().cpu().numpy() # Shape: [H, W]
+
+    # Threshold probabilities to get binary mask (0 or 1)
+    binary_mask_np = (mask_pred_np > 0.5).astype(np.uint8)
+
+    # Convert binary mask to a displayable format (e.g., 0 or 255)
+    display_mask = (binary_mask_np * 255) # Shape: [H, W]
+
+    # Resize mask back to original image size for overlay (optional, better overlay quality)
+    orig_h, orig_w = input_image_rgb.shape[:2]
+    display_mask_resized = cv2.resize(display_mask, (orig_w, orig_h), interpolation=cv2.INTER_NEAREST)
+
+    # 4. Create Overlay
+    # Convert single-channel mask to 3 channels to overlay on RGB image
+    mask_rgb = cv2.cvtColor(display_mask_resized, cv2.COLOR_GRAY2RGB)
+    # Make the mask red where segmentation is present
+    mask_rgb[:, :, 0] = 0 # Zero out Blue channel
+    mask_rgb[:, :, 1] = 0 # Zero out Green channel
+    # Where mask_rgb is red (255), keep original image pixel, otherwise blend
+    overlay_image = cv2.addWeighted(input_image_rgb, 0.7, mask_rgb, 0.3, 0)
+    # Highlight only the segmented area more distinctly
+    highlighted_area = cv2.bitwise_and(input_image_rgb, input_image_rgb, mask=display_mask_resized)
+    overlay_image = cv2.addWeighted(input_image_rgb, 0.7, highlighted_area, 0.9, 0) # Blend original with highlighted
+
+    # Return original, mask (resized), overlay
+    # Gradio expects NumPy arrays
+    #return input_image_rgb, display_mask_resized, overlay_image
+    return display_mask_resized, overlay_image
+
+
+# --- Gradio Interface ---
+print("Launching Gradio Interface...")
+
+with gr.Blocks() as demo:
+    gr.Markdown("# Dermatology Image Segmentation (UNet ResNet34)")
+    gr.Markdown("Upload a dermatology image to see the predicted segmentation mask using a trained UNet model.")
+    
+    with gr.Row():            # Creates a horizontal container
+        inp = gr.Image(type="numpy", label="Input Image")
+        out_mask = gr.Image(type="numpy", label="Segmentation Mask")
+        out_overlay = gr.Image(type="numpy", label="Overlay")
+    
+    # Hook up the function
+    inp.change(fn=segment_image, inputs=inp, outputs=[out_mask, out_overlay])
+    
+    # (Optional) add example images
+    # gr.Examples(examples=[["examples/img1.jpg"], ["examples/img2.jpg"]], 
+    #             inputs=inp, outputs=[out_mask, out_overlay])
+    
+    # Disable flagging
+    
+if __name__ == "__main__":
+    demo.launch(share=True) # Share=True to create public link

requirements.txt

@@ -0,0 +1,104 @@
+absl-py==2.2.2
+aiofiles==24.1.0
+aiohappyeyeballs==2.6.1
+aiohttp==3.11.18
+aiosignal==1.3.2
+albucore==0.0.23
+albumentations==2.0.5
+annotated-types==0.7.0
+anyio==4.9.0
+attrs==25.3.0
+certifi==2025.4.26
+charset-normalizer==3.4.1
+click==8.1.8
+fastapi==0.115.12
+ffmpy==0.5.0
+filelock==3.18.0
+frozenlist==1.6.0
+fsspec==2025.3.2
+gradio==5.27.0
+gradio_client==1.9.0
+groovy==0.1.2
+grpcio==1.71.0
+h11==0.16.0
+httpcore==1.0.9
+httpx==0.28.1
+huggingface-hub==0.30.2
+idna==3.10
+Jinja2==3.1.6
+lightning-utilities==0.14.3
+Markdown==3.8
+markdown-it-py==3.0.0
+MarkupSafe==3.0.2
+mdurl==0.1.2
+mpmath==1.3.0
+multidict==6.4.3
+networkx==3.4.2
+numpy==2.2.5
+nvidia-cublas-cu12==12.6.4.1
+nvidia-cuda-cupti-cu12==12.6.80
+nvidia-cuda-nvrtc-cu12==12.6.77
+nvidia-cuda-runtime-cu12==12.6.77
+nvidia-cudnn-cu12==9.5.1.17
+nvidia-cufft-cu12==11.3.0.4
+nvidia-cufile-cu12==1.11.1.6
+nvidia-curand-cu12==10.3.7.77
+nvidia-cusolver-cu12==11.7.1.2
+nvidia-cusparse-cu12==12.5.4.2
+nvidia-cusparselt-cu12==0.6.3
+nvidia-nccl-cu12==2.26.2
+nvidia-nvjitlink-cu12==12.6.85
+nvidia-nvtx-cu12==12.6.77
+opencv-python==4.11.0.86
+opencv-python-headless==4.11.0.86
+orjson==3.10.16
+packaging==25.0
+pandas==2.2.3
+pillow==11.2.1
+propcache==0.3.1
+protobuf==6.30.2
+pydantic==2.11.3
+pydantic_core==2.33.1
+pydub==0.25.1
+Pygments==2.19.1
+python-dateutil==2.9.0.post0
+python-multipart==0.0.20
+pytorch-lightning==2.5.1.post0
+pytz==2025.2
+PyYAML==6.0.2
+requests==2.32.3
+rich==14.0.0
+ruff==0.11.7
+safehttpx==0.1.6
+safetensors==0.5.3
+scipy==1.15.2
+segmentation_models_pytorch==0.5.0
+semantic-version==2.10.0
+setuptools==75.8.0
+shellingham==1.5.4
+simsimd==6.2.1
+six==1.17.0
+sniffio==1.3.1
+starlette==0.46.2
+stringzilla==3.12.5
+sympy==1.14.0
+tensorboard==2.19.0
+tensorboard-data-server==0.7.2
+timm==1.0.15
+tomlkit==0.13.2
+torch==2.7.0
+torchaudio==2.7.0
+torchmetrics==1.7.1
+torchvision==0.22.0
+tqdm==4.67.1
+triton==3.3.0
+typer==0.15.2
+typing-inspection==0.4.0
+typing_extensions==4.13.2
+tzdata==2025.2
+urllib3==2.4.0
+uvicorn==0.34.2
+websockets==15.0.1
+Werkzeug==3.1.3
+wheel==0.45.1
+yarl==1.20.0

unet_derm_final_model.pth

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