Upload 21 files

.gitignore

@@ -0,0 +1,9 @@
+# Ignore virtual environment
+venv/
+__pycache__/
+
+# Ignore model weights (GitHub has a 100MB limit)
+*.pth
+
+# Ignore system files
+.DS_Store

.idea/.gitignore

@@ -0,0 +1,3 @@
+# 默认忽略的文件
+/shelf/
+/workspace.xml

.idea/image.iml

@@ -0,0 +1,8 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<module type="PYTHON_MODULE" version="4">
+  <component name="NewModuleRootManager">
+    <content url="file://$MODULE_DIR$" />
+    <orderEntry type="jdk" jdkName="Python 3.9 (look)" jdkType="Python SDK" />
+    <orderEntry type="sourceFolder" forTests="false" />
+  </component>
+</module>

.idea/inspectionProfiles/Project_Default.xml

@@ -0,0 +1,77 @@
+<component name="InspectionProjectProfileManager">
+  <profile version="1.0">
+    <option name="myName" value="Project Default" />
+    <inspection_tool class="PyCompatibilityInspection" enabled="true" level="WARNING" enabled_by_default="true">
+      <option name="ourVersions">
+        <value>
+          <list size="3">
+            <item index="0" class="java.lang.String" itemvalue="2.7" />
+            <item index="1" class="java.lang.String" itemvalue="3.13" />
+            <item index="2" class="java.lang.String" itemvalue="3.9" />
+          </list>
+        </value>
+      </option>
+    </inspection_tool>
+    <inspection_tool class="PyPackageRequirementsInspection" enabled="true" level="WARNING" enabled_by_default="true">
+      <option name="ignoredPackages">
+        <value>
+          <list size="53">
+            <item index="0" class="java.lang.String" itemvalue="tensorboard" />
+            <item index="1" class="java.lang.String" itemvalue="timm" />
+            <item index="2" class="java.lang.String" itemvalue="absl-py" />
+            <item index="3" class="java.lang.String" itemvalue="shapely" />
+            <item index="4" class="java.lang.String" itemvalue="joblib" />
+            <item index="5" class="java.lang.String" itemvalue="threadpoolctl" />
+            <item index="6" class="java.lang.String" itemvalue="huggingface-hub" />
+            <item index="7" class="java.lang.String" itemvalue="PyYAML" />
+            <item index="8" class="java.lang.String" itemvalue="setuptools" />
+            <item index="9" class="java.lang.String" itemvalue="fsspec" />
+            <item index="10" class="java.lang.String" itemvalue="filelock" />
+            <item index="11" class="java.lang.String" itemvalue="qudida" />
+            <item index="12" class="java.lang.String" itemvalue="pip" />
+            <item index="13" class="java.lang.String" itemvalue="certifi" />
+            <item index="14" class="java.lang.String" itemvalue="warmup-scheduler" />
+            <item index="15" class="java.lang.String" itemvalue="portalocker" />
+            <item index="16" class="java.lang.String" itemvalue="nibabel" />
+            <item index="17" class="java.lang.String" itemvalue="h5py" />
+            <item index="18" class="java.lang.String" itemvalue="loguru" />
+            <item index="19" class="java.lang.String" itemvalue="contourpy" />
+            <item index="20" class="java.lang.String" itemvalue="fonttools" />
+            <item index="21" class="java.lang.String" itemvalue="imageio" />
+            <item index="22" class="java.lang.String" itemvalue="fairscale" />
+            <item index="23" class="java.lang.String" itemvalue="matplotlib" />
+            <item index="24" class="java.lang.String" itemvalue="charset-normalizer" />
+            <item index="25" class="java.lang.String" itemvalue="MedPy" />
+            <item index="26" class="java.lang.String" itemvalue="SimpleITK" />
+            <item index="27" class="java.lang.String" itemvalue="idna" />
+            <item index="28" class="java.lang.String" itemvalue="scikit-image" />
+            <item index="29" class="java.lang.String" itemvalue="yapf" />
+            <item index="30" class="java.lang.String" itemvalue="numpy" />
+            <item index="31" class="java.lang.String" itemvalue="requests" />
+            <item index="32" class="java.lang.String" itemvalue="opencv-python-headless" />
+            <item index="33" class="java.lang.String" itemvalue="seaborn" />
+            <item index="34" class="java.lang.String" itemvalue="pthflops" />
+            <item index="35" class="java.lang.String" itemvalue="PyWavelets" />
+            <item index="36" class="java.lang.String" itemvalue="zipp" />
+            <item index="37" class="java.lang.String" itemvalue="lazy_loader" />
+            <item index="38" class="java.lang.String" itemvalue="scipy" />
+            <item index="39" class="java.lang.String" itemvalue="opencv-python" />
+            <item index="40" class="java.lang.String" itemvalue="wheel" />
+            <item index="41" class="java.lang.String" itemvalue="packaging" />
+            <item index="42" class="java.lang.String" itemvalue="addict" />
+            <item index="43" class="java.lang.String" itemvalue="albumentations" />
+            <item index="44" class="java.lang.String" itemvalue="termcolor" />
+            <item index="45" class="java.lang.String" itemvalue="importlib-resources" />
+            <item index="46" class="java.lang.String" itemvalue="typing_extensions" />
+            <item index="47" class="java.lang.String" itemvalue="pytz" />
+            <item index="48" class="java.lang.String" itemvalue="einops" />
+            <item index="49" class="java.lang.String" itemvalue="Pillow" />
+            <item index="50" class="java.lang.String" itemvalue="torch" />
+            <item index="51" class="java.lang.String" itemvalue="torchvision" />
+            <item index="52" class="java.lang.String" itemvalue="torchsummary" />
+          </list>
+        </value>
+      </option>
+    </inspection_tool>
+  </profile>
+</component>

.idea/inspectionProfiles/profiles_settings.xml

@@ -0,0 +1,6 @@
+<component name="InspectionProjectProfileManager">
+  <settings>
+    <option name="USE_PROJECT_PROFILE" value="false" />
+    <version value="1.0" />
+  </settings>
+</component>

.idea/misc.xml

@@ -0,0 +1,7 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<project version="4">
+  <component name="Black">
+    <option name="sdkName" value="Python 3.9 (look)" />
+  </component>
+  <component name="ProjectRootManager" version="2" project-jdk-name="Python 3.9 (look)" project-jdk-type="Python SDK" />
+</project>

.idea/modules.xml

@@ -0,0 +1,8 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<project version="4">
+  <component name="ProjectModuleManager">
+    <modules>
+      <module fileurl="file://$PROJECT_DIR$/.idea/image.iml" filepath="$PROJECT_DIR$/.idea/image.iml" />
+    </modules>
+  </component>
+</project>

.idea/workspace.xml

@@ -0,0 +1,49 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<project version="4">
+  <component name="AutoImportSettings">
+    <option name="autoReloadType" value="SELECTIVE" />
+  </component>
+  <component name="ChangeListManager">
+    <list default="true" id="2992db4b-be05-4abd-b057-113b2d800a0e" name="更改" comment="" />
+    <option name="SHOW_DIALOG" value="false" />
+    <option name="HIGHLIGHT_CONFLICTS" value="true" />
+    <option name="HIGHLIGHT_NON_ACTIVE_CHANGELIST" value="false" />
+    <option name="LAST_RESOLUTION" value="IGNORE" />
+  </component>
+  <component name="MarkdownSettingsMigration">
+    <option name="stateVersion" value="1" />
+  </component>
+  <component name="ProjectColorInfo"><![CDATA[{
+  "associatedIndex": 1
+}]]></component>
+  <component name="ProjectId" id="2y2VrJeUMs6MgMG5dRwdH89cBQE" />
+  <component name="ProjectViewState">
+    <option name="hideEmptyMiddlePackages" value="true" />
+    <option name="showLibraryContents" value="true" />
+  </component>
+  <component name="PropertiesComponent"><![CDATA[{
+  "keyToString": {
+    "RunOnceActivity.OpenProjectViewOnStart": "true",
+    "RunOnceActivity.ShowReadmeOnStart": "true",
+    "last_opened_file_path": "E:/挑战杯/image"
+  }
+}]]></component>
+  <component name="SharedIndexes">
+    <attachedChunks>
+      <set>
+        <option value="bundled-python-sdk-09665e90c3a7-d3b881c8e49f-com.jetbrains.pycharm.community.sharedIndexes.bundled-PC-233.15026.15" />
+      </set>
+    </attachedChunks>
+  </component>
+  <component name="SpellCheckerSettings" RuntimeDictionaries="0" Folders="0" CustomDictionaries="0" DefaultDictionary="应用程序级" UseSingleDictionary="true" transferred="true" />
+  <component name="TaskManager">
+    <task active="true" id="Default" summary="默认任务">
+      <changelist id="2992db4b-be05-4abd-b057-113b2d800a0e" name="更改" comment="" />
+      <created>1749032852317</created>
+      <option name="number" value="Default" />
+      <option name="presentableId" value="Default" />
+      <updated>1749032852317</updated>
+    </task>
+    <servers />
+  </component>
+</project>

Dockerfile

@@ -0,0 +1,24 @@
+# Use the official Python image as a base
+FROM python:3.9-slim
+
+# Install system dependencies for OpenCV and wget
+RUN apt-get update && apt-get install -y libgl1-mesa-glx libglib2.0-0 wget && rm -rf /var/lib/apt/lists/*
+
+# Set working directory inside the container
+WORKDIR /app
+
+# Copy and install dependencies
+COPY requirements.txt .
+RUN pip install --no-cache-dir -r requirements.txt
+
+# Download model weights from Hugging Face
+RUN wget -O best_swin_upernet_main.pth "https://huggingface.co/samyakshrestha/swin-medical-segmentation/resolve/main/best_swin_upernet_main.pth?download=true"
+
+# Copy the rest of the app code
+COPY . .
+
+# Expose the port
+EXPOSE 8000
+
+# Run FastAPI with uvicorn
+CMD ["uvicorn", "main:app", "--host", "0.0.0.0", "--port", "8000"]

README.md

@@ -1,12 +1,89 @@
----
-title: Fenge1
-emoji: 👀
-colorFrom: yellow
-colorTo: pink
-sdk: gradio
-sdk_version: 5.32.1
-app_file: app.py
-pinned: false
----
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# Skin Lesion Segmentation API with Uncertainty Estimation
+
+This repository contains a production-ready FastAPI service for performing semantic segmentation of skin lesions, trained on the ISIC 2018 dataset. The segmentation model uses a Swin-Tiny transformer encoder with a custom UPerNet decoder, and incorporates Monte Carlo (MC) Dropout at inference time to estimate predictive uncertainty.
+
+## Features
+
+- RESTful FastAPI endpoint for real-time image segmentation
+- Swin-Tiny + UPerNet architecture for accurate lesion boundary detection
+- Monte Carlo Dropout-based uncertainty quantification
+- Heatmap visualization of model uncertainty using OpenCV colormaps
+- Side-by-side output of the predicted segmentation mask and uncertainty map
+- Clean, modular codebase suitable for further extension and deployment
+
+## Inference Workflow
+
+1. The API accepts a dermoscopic image uploaded via a POST request to the `/predict/` endpoint.
+2. The image is resized, normalized, and passed through the segmentation model.
+3. The model performs multiple stochastic forward passes using MC Dropout.
+4. The mean prediction is binarized to produce a segmentation mask.
+5. The standard deviation across predictions is visualized as a heatmap to express uncertainty.
+6. The API returns both the binary mask and uncertainty heatmap as a combined PNG image.
+
+## File Structure
+
+| File              | Description                                                  |
+|-------------------|--------------------------------------------------------------|
+| `main.py`         | FastAPI application and endpoint logic                       |
+| `model.py`        | Model architecture and MC Dropout inference utilities        |
+| `requirements.txt`| Required Python packages                                     |
+| `.gitignore`      | Specifies ignored files (e.g., `venv/`, model weights)       |
+
+## Requirements
+
+- Python 3.9+
+- PyTorch
+- FastAPI
+- OpenCV
+- Pillow
+- timm
+- torchvision
+- uvicorn
+
+Install dependencies:
+
+```bash
+pip install -r requirements.txt
+```
+
+
+## Model Weights
+
+To run the API, you will need to download the pretrained model weights and place the file in the root directory of this repository.
+
+**Download link:**
+
+[best_swin_upernet_main.pth](https://huggingface.co/samyakshrestha/swin-medical-segmentation/resolve/main/best_swin_upernet_main.pth?download=true)
+
+After downloading, ensure the file is named exactly "best_swin_upernet_main.pth" and saved in the same directory as `main.py` and `model.py`.
+
+This weight file contains the Swin-Tiny UPerNet model trained on the ISIC 2018 skin lesion segmentation dataset.
+
+
+## Dataset
+
+This project uses the ISIC 2018 Challenge Dataset for training, which contains dermoscopic images of skin lesions annotated for binary segmentation.
+
+Dataset access: https://challenge.isic-archive.com/data/
+
+## Docker Support
+
+This project includes a `Dockerfile` for containerized deployment. Run the API with Docker:
+
+### Build the Docker image
+```bash
+docker build -t medical-segmentation-api .
+```
+
+### Run the container
+```bash
+docker run -p 8000:8000 medical-segmentation-api
+```
+
+
+The API will then be available at:
+http://localhost:8000/predict/
+
+## Author
+
+Samyak Shrestha 

aa.py

@@ -0,0 +1,9 @@
+import requests
+
+url = "http://127.0.0.1:4011/predict/"
+files = {'file': open('D:\zkwg\image\data\iamge1.jpg', 'rb')}  # 替换为您的图像路径
+response = requests.post(url, files=files)
+
+# 保存返回的预测结果
+with open('D:\zkwg\image\data\prediction_result.png', 'wb') as f:
+    f.write(response.content)

app.py

@@ -0,0 +1,128 @@
+import streamlit as st
+import torch
+from PIL import Image
+import numpy as np
+import cv2
+from model import load_model, predict_with_uncertainty
+import torchvision.transforms as transforms
+from transformers import AutoProcessor, AutoModelForImageTextToText
+from io import BytesIO
+
+# 设置页面配置
+st.set_page_config(
+    page_title="医学图像分析系统",
+    page_icon="🏥",
+    layout="wide"
+)
+
+# 初始化模型
+@st.cache_resource
+def load_models():
+    # 加载分割模型
+    seg_model = load_model()
+    seg_model.eval()
+    
+    # 加载分析模型
+    model_id = "google/medgemma-4b-it"
+    analysis_model = AutoModelForImageTextToText.from_pretrained(
+        model_id,
+        torch_dtype=torch.bfloat16,
+        device_map="auto",
+        token="HUGGINGFACE_TOKEN"
+    )
+    processor = AutoProcessor.from_pretrained(
+        model_id,
+        token="HUGGINGFACE_TOKEN"
+    )
+    
+    return seg_model, analysis_model, processor
+
+# 页面标题
+st.title("🏥 医学图像分析系统")
+st.markdown("---")
+
+# 加载模型
+with st.spinner("正在加载模型..."):
+    seg_model, analysis_model, processor = load_models()
+
+# 创建两列布局
+col1, col2 = st.columns(2)
+
+with col1:
+    st.subheader("📤 上传图片")
+    uploaded_file = st.file_uploader("选择一张图片", type=["png", "jpg", "jpeg"])
+    
+    if uploaded_file is not None:
+        # 显示原始图片
+        image = Image.open(uploaded_file).convert("RGB")
+        st.image(image, caption="原始图片", use_column_width=True)
+        
+        # 处理图片
+        if st.button("开始分析"):
+            with st.spinner("正在处理..."):
+                # 图像分割
+                image_resized = image.resize((224, 224))
+                transform = transforms.ToTensor()
+                image_tensor = transform(image_resized).unsqueeze(0)
+                
+                # 执行分割
+                preds_mean, preds_uncertainty = predict_with_uncertainty(image_tensor)
+                
+                # 生成分割掩码
+                pred_binary = (preds_mean > 0.5).astype(np.uint8) * 255
+                mask_image = Image.fromarray(pred_binary).convert("L")
+                
+                # 处理不确定性图
+                uncertainty = (preds_uncertainty - preds_uncertainty.min()) / (preds_uncertainty.max() - preds_uncertainty.min() + 1e-8)
+                uncertainty_colormap = cv2.applyColorMap((uncertainty * 255).astype(np.uint8), cv2.COLORMAP_INFERNO)
+                uncertainty_image = Image.fromarray(uncertainty_colormap).convert("RGB")
+                
+                # 合并图片
+                combined = Image.new("RGB", (mask_image.width + uncertainty_image.width, mask_image.height))
+                combined.paste(mask_image.convert("RGB"), (0, 0))
+                combined.paste(uncertainty_image, (mask_image.width, 0))
+                
+                # 图像分析
+                messages = [
+                    {
+                        "role": "system",
+                        "content": [{"type": "text", "text": "你是一名皮肤病专家，请使用中文分析图片."}]
+                    },
+                    {
+                        "role": "user",
+                        "content": [
+                            {"type": "text", "text": "这是一张皮肤病的图片，帮我分析一下"},
+                            {"type": "image", "image": image}
+                        ]
+                    }
+                ]
+                
+                inputs = processor.apply_chat_template(
+                    messages, add_generation_prompt=True, tokenize=True,
+                    return_dict=True, return_tensors="pt"
+                ).to(analysis_model.device, dtype=torch.bfloat16)
+                
+                input_len = inputs["input_ids"].shape[-1]
+                
+                with torch.inference_mode():
+                    generation = analysis_model.generate(**inputs, max_new_tokens=200, do_sample=False)
+                    generation = generation[0][input_len:]
+                
+                analysis_text = processor.decode(generation, skip_special_tokens=True)
+                
+                # 显示结果
+                with col2:
+                    st.subheader("📊 分析结果")
+                    st.image(combined, caption="分割结果", use_column_width=True)
+                    st.markdown("### 📝 图像分析")
+                    st.write(analysis_text)
+
+# 添加页脚
+st.markdown("---")
+st.markdown("### 使用说明")
+st.markdown("""
+1. 在左侧上传一张医学图像
+2. 点击"开始分析"按钮
+3. 系统将自动进行图像分割和分析
+4. 右侧将显示分割结果和分析报告
+""") 

main.py

@@ -0,0 +1,60 @@
+import logging
+from fastapi import FastAPI, UploadFile, File
+from fastapi.responses import Response
+import uvicorn
+from model import load_model, predict_with_uncertainty
+import torchvision.transforms as transforms
+from PIL import Image
+from io import BytesIO
+import numpy as np
+import cv2
+import torch
+
+app = FastAPI()
+
+# Load the model when the API starts
+model = load_model()
+model.eval()
+
+def convert_to_image(array, colormap=None):
+    array = (array * 255).astype(np.uint8)  # Normalize to 0–255 range
+    if colormap is not None:
+        array = cv2.applyColorMap(array, colormap)
+    return Image.fromarray(array)
+
+@app.post("/predict/")
+async def predict_mask(file: UploadFile = File(...)):
+    # Read and preprocess the image
+    image = Image.open(BytesIO(await file.read())).convert("RGB")
+    image = image.resize((224, 224))
+    transform = transforms.ToTensor()
+    image_tensor = transform(image).unsqueeze(0)
+
+    # Perform MC Dropout Inference
+    preds_mean, preds_uncertainty = predict_with_uncertainty(image_tensor)
+
+    # Binary mask (0 or 255)
+    pred_binary = (preds_mean > 0.5).astype(np.uint8) * 255
+    mask_image = Image.fromarray(pred_binary).convert("L")
+
+    # Normalize and apply colormap to uncertainty
+    uncertainty = (preds_uncertainty - preds_uncertainty.min()) / (preds_uncertainty.max() - preds_uncertainty.min() + 1e-8)
+    uncertainty_colormap = cv2.applyColorMap((uncertainty * 255).astype(np.uint8), cv2.COLORMAP_INFERNO)
+    uncertainty_image = Image.fromarray(uncertainty_colormap).convert("RGB")
+
+    # Combine side by side
+    combined = Image.new("RGB", (mask_image.width + uncertainty_image.width, mask_image.height))
+    combined.paste(mask_image.convert("RGB"), (0, 0))
+    combined.paste(uncertainty_image, (mask_image.width, 0))
+
+    # Save to buffer and return
+    img_io = BytesIO()
+    combined.save(img_io, format="PNG")
+    img_io.seek(0)
+
+    return Response(content=img_io.getvalue(), media_type="image/png")
+
+if __name__ == "__main__":
+    # 直接启动模式，适合各种操作系统环境
+    logging.info("启动MCP客户端API服务...")
+    uvicorn.run(app, host="0.0.0.0", port=4011)

model.py

@@ -0,0 +1,130 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import timm
+import numpy as np
+import cv2
+
+# -------------------------------
+# Define Pyramid Pooling Module (with GroupNorm)
+# -------------------------------
+class PyramidPoolingModule(nn.Module):
+    def __init__(self, in_channels, pool_sizes=[1, 2, 3, 6]):
+        super().__init__()
+        self.pool_layers = nn.ModuleList([
+            nn.Sequential(
+                nn.AdaptiveAvgPool2d(pool_size),
+                nn.Conv2d(in_channels, in_channels // 4, kernel_size=1, bias=False),
+                nn.GroupNorm(num_groups=8, num_channels=in_channels // 4),
+                nn.ReLU(inplace=True)
+            ) for pool_size in pool_sizes
+        ])
+        total_channels = in_channels + len(pool_sizes) * (in_channels // 4)
+        self.conv = nn.Conv2d(total_channels, in_channels, kernel_size=1, bias=False)
+
+    def forward(self, x):
+        pooled_features = [x]
+        for layer in self.pool_layers:
+            pooled = layer(x)
+            pooled = F.interpolate(pooled, size=x.shape[2:], mode='bilinear', align_corners=False)
+            pooled_features.append(pooled)
+        x = torch.cat(pooled_features, dim=1)
+        x = self.conv(x)
+        return x
+
+# -------------------------------
+# Define UPerNet Decoder (With Dropout)
+# -------------------------------
+class UPerNetDecoder(nn.Module):
+    def __init__(self, encoder_channels, num_classes=1, dropout_rate=0.1):
+        super().__init__()
+        self.ppm = PyramidPoolingModule(encoder_channels[-1])
+        self.lateral_conv2 = nn.Conv2d(encoder_channels[2], encoder_channels[-1], kernel_size=1)
+        self.conv3 = nn.Sequential(
+            nn.Conv2d(encoder_channels[-1], encoder_channels[2], kernel_size=1),
+            nn.Dropout2d(p=dropout_rate)
+        )
+        self.lateral_conv1 = nn.Conv2d(encoder_channels[1], encoder_channels[2], kernel_size=1)
+        self.conv2 = nn.Sequential(
+            nn.Conv2d(encoder_channels[2], encoder_channels[1], kernel_size=1),
+            nn.Dropout2d(p=dropout_rate)
+        )
+        self.lateral_conv0 = nn.Conv2d(encoder_channels[0], encoder_channels[1], kernel_size=1)
+        self.conv1 = nn.Sequential(
+            nn.Conv2d(encoder_channels[1], encoder_channels[0], kernel_size=1),
+            nn.Dropout2d(p=dropout_rate)
+        )
+        self.segmentation_head = nn.Conv2d(encoder_channels[0], num_classes, kernel_size=1)
+
+    def forward(self, features):
+        f0, f1, f2, f3 = features
+        x3 = self.ppm(f3)
+        x3_up = F.interpolate(x3, size=f2.shape[2:], mode="bilinear", align_corners=False)
+        fuse2 = x3_up + self.lateral_conv2(f2)
+        fuse2 = self.conv3(fuse2)
+        fuse2_up = F.interpolate(fuse2, size=f1.shape[2:], mode="bilinear", align_corners=False)
+        fuse1 = fuse2_up + self.lateral_conv1(f1)
+        fuse1 = self.conv2(fuse1)
+        fuse1_up = F.interpolate(fuse1, size=f0.shape[2:], mode="bilinear", align_corners=False)
+        fuse0 = fuse1_up + self.lateral_conv0(f0)
+        fuse0 = self.conv1(fuse0)
+        x_out = F.interpolate(fuse0, size=(224, 224), mode="bilinear", align_corners=False)
+        output = self.segmentation_head(x_out)
+        return output
+
+# -------------------------------
+# Define Swin-Tiny UPerNet Model
+# -------------------------------
+class SwinTinyUPerNet(nn.Module):
+    def __init__(self, num_classes=1, dropout_rate=0.1):
+        super().__init__()
+        self.encoder = timm.create_model(
+            "swin_tiny_patch4_window7_224.ms_in22k_ft_in1k",
+            pretrained=True,
+            features_only=True
+        )
+        encoder_channels = self.encoder.feature_info.channels()
+        self.decoder = UPerNetDecoder(encoder_channels, num_classes, dropout_rate=dropout_rate)
+
+    def forward(self, x):
+        features = self.encoder(x)
+        features = [f.permute(0, 3, 1, 2) if f.dim() == 4 else f for f in features]
+        output = self.decoder(features)
+        return F.interpolate(output, size=(224, 224), mode="bilinear", align_corners=False)
+
+# -------------------------------
+# Load the Model
+# -------------------------------
+def load_model():
+    model = SwinTinyUPerNet(num_classes=1)
+    model.load_state_dict(torch.load("best_swin_upernet_main.pth", map_location=torch.device("cpu")), strict=False)
+    model.eval()
+    return model
+
+# -------------------------------
+# Enable Dropout at Inference Time
+# -------------------------------
+def enable_dropout(m):
+    if isinstance(m, nn.Dropout) or isinstance(m, nn.Dropout2d):
+        m.train()
+
+# -------------------------------
+# Perform Inference with MC Dropout
+# -------------------------------
+def predict_with_uncertainty(image_tensor, num_samples=10):
+    model = load_model()
+    model.apply(enable_dropout)
+    preds_list = []
+
+    with torch.no_grad():
+        for _ in range(num_samples):
+            preds = torch.sigmoid(model(image_tensor))
+            preds_list.append(preds)
+
+    preds_array = torch.stack(preds_list, dim=0)
+    preds_mean = preds_array.mean(dim=0).squeeze().cpu().numpy()
+    preds_uncertainty = preds_array.std(dim=0).squeeze().cpu().numpy()
+
+    # Normalize uncertainty map
+    preds_uncertainty = (preds_uncertainty - preds_uncertainty.min()) / (preds_uncertainty.max() - preds_uncertainty.min() + 1e-8)
+    return preds_mean, preds_uncertainty

requirements.txt

@@ -0,0 +1,15 @@
+streamlit==1.32.0
+torch==2.2.0
+torchvision==0.17.0
+Pillow==10.2.0
+numpy==1.26.4
+opencv-python==4.9.0.80
+fastapi==0.110.0
+uvicorn==0.27.1
+python-multipart==0.0.9
+huggingface-hub>=0.30.0
+accelerate>=0.27.0
+sentencepiece==0.2.0
+protobuf==4.25.3
+einops==0.7.0
+safetensors==0.4.2