Mufasa06/WoundClassification

1import os2import sys3import yaml4import torch5import numpy as np6import pandas as pd7from PIL import Image8import torchvision.transforms as T9import streamlit as st10 11# Ensure src is in PATH12sys.path.append(os.path.join(os.path.dirname(__file__), 'src'))13 14from single_inference_plot import AttentionMaskGradCAM, tensor_to_img, overlay_cam_mask, overlay_mask15from train_segmentation import AttentionUNetEfficientNetB016from train_classification import SoftGuidedAttentionEfficientNet17 18# Page configuration19st.set_page_config(page_title="AI Wound Classifier", layout="wide")20 21# Custom CSS for styling22st.markdown("""23<style>24.severity-high { padding:10px; border-left: 5px solid #28a745; background: #e8f5e9; border-radius: 4px; }25.severity-moderate { padding:10px; border-left: 5px solid #ffc107; background: #fff8e1; border-radius: 4px; }26.severity-low { padding:10px; border-left: 5px solid #dc3545; background: #ffebee; border-radius: 4px; }27.report-box { padding:20px; background: #f8f9fa; border-radius: 8px; border: 1px solid #dee2e6; margin-top: 10px; }28</style>29""", unsafe_allow_html=True)30 31# Helper function to cache model loading32@st.cache_resource33def load_models():34    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")35    36    config_path = "src/config.yaml"37    with open(config_path, "r") as f:38        config = yaml.safe_load(f)39 40    cfg_seg = config['segmentation']41    cfg_cls = config['classification']42    cfg_inf = config['inference']43 44    # 1. Load Segmentation Model45    try:46        seg_model = torch.jit.load(cfg_inf['seg_model_path'], map_location=device)47        seg_model.eval()48    except Exception:49        seg_model = AttentionUNetEfficientNetB0(50            in_channels=cfg_seg['in_channels'],51            out_channels=cfg_seg['out_channels'],52            pretrained=False53        ).to(device)54        seg_model.load_state_dict(torch.load(cfg_inf['seg_model_path'], map_location=device))55        seg_model.eval()56 57    # 2. Load Classification Model58    model_gradcam_path = "models/classification_runs/model_gradcam.pth"59    cls_model = SoftGuidedAttentionEfficientNet(num_classes=4).to(device)60    cls_model.load_state_dict(torch.load(model_gradcam_path, map_location=device))61    cls_model.eval()62 63    # Transforms64    seg_img_size = tuple(cfg_seg['img_size'])65    resize_to_tensor_seg = T.Compose([66        T.Resize(seg_img_size, interpolation=T.InterpolationMode.BILINEAR),67        T.ToTensor(),68    ])69    normalize_seg = T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])70 71    cls_img_size = tuple(cfg_cls['img_size'])72    resize_to_tensor_cls = T.Compose([73        T.Resize(cls_img_size, interpolation=T.InterpolationMode.BILINEAR),74        T.ToTensor(),75    ])76    77    return device, config, seg_model, cls_model, resize_to_tensor_seg, normalize_seg, resize_to_tensor_cls78 79device, config, seg_model, cls_model, resize_to_tensor_seg, normalize_seg, resize_to_tensor_cls = load_models()80cfg_inf = config['inference']81seg_img_size = tuple(config['segmentation']['img_size'])82 83class_names = ["diabetic", "masd", "pressure", "venous"]84friendly_names = {85    "diabetic": "Diabetic Foot Ulcer",86    "masd": "Moisture Associated Skin Damage (MASD)",87    "pressure": "Pressure Ulcer",88    "venous": "Venous Leg Ulcer"89}90 91def analyze_wound(img_input):92    if isinstance(img_input, np.ndarray):93        img_pil = Image.fromarray(img_input).convert("RGB")94    else:95        img_pil = img_input.convert("RGB")96 97    with torch.no_grad():98        # --- SEGMENTATION ---99        x_resized = resize_to_tensor_seg(img_pil)100        x_norm = normalize_seg(x_resized).unsqueeze(0).to(device)101        102        seg_pred = seg_model(x_norm)103        if isinstance(seg_pred, (tuple, list)): 104            seg_pred = seg_pred[0]105            106        seg_pred = torch.sigmoid(seg_pred)107 108        if seg_pred.ndim == 4 and seg_pred.shape[1] == 1:109            seg_probs = seg_pred[0, 0].detach().cpu().numpy()110        else:111            seg_probs = seg_pred.squeeze().detach().cpu().numpy()112            113        mask_np = (seg_probs >= cfg_inf['threshold']).astype(np.uint8)114        115        mask_pil = Image.fromarray((mask_np * 255).astype(np.uint8))116        mask_pil = mask_pil.resize(seg_img_size, resample=Image.NEAREST)117 118        # --- CLASSIFICATION ---119        rgb_pil_saved = tensor_to_img(x_resized.unsqueeze(0))120        rgb_pil_to_save = Image.fromarray((rgb_pil_saved * 255.0).astype(np.uint8))121        122        img_cls = resize_to_tensor_cls(rgb_pil_to_save).unsqueeze(0).to(device)123        124        mask_tensor_cls = T.ToTensor()(mask_pil).unsqueeze(0).to(device)125        mask_tensor_cls = (mask_tensor_cls > 0.5).float()126 127        cls_logits = cls_model(img_cls, mask_tensor_cls)128        # Note: if it returns a tuple, grab the first element129        if isinstance(cls_logits, (tuple, list)):130            cls_logits = cls_logits[0]131            132        probs = torch.softmax(cls_logits, 1)[0].detach().cpu().numpy()133        predicted_class = int(torch.argmax(cls_logits, dim=1).item())134 135    pred_name = class_names[predicted_class]136    confidence = float(probs[predicted_class])137    138    # Format probabilities139    confidences_dict = {friendly_names[name]: float(prob) for name, prob in zip(class_names, probs)}140 141    # --- GRADCAM ---142    gradcam = AttentionMaskGradCAM(cls_model, use_ts=False)143    cam = gradcam(img_cls, mask_tensor_cls, predicted_class)144    gradcam.remove()145 146    # --- VISUALIZATION ---147    img_np = tensor_to_img(img_cls)148    mask_np_cls = mask_tensor_cls[0, 0].cpu().numpy()149    150    # 1. Raw mask image151    mask_display = Image.fromarray((mask_np_cls * 255).astype(np.uint8)).convert("RGB")152    # 2. Mask overlay on original image153    mask_overlay_img = overlay_mask(img_np, mask_np_cls)154    mask_overlay_display = Image.fromarray((mask_overlay_img * 255).astype(np.uint8))155    # 3. Combined GradCAM + Mask overlay156    combined = overlay_cam_mask(img_np, cam, mask_np_cls)157    combined_display = Image.fromarray((combined * 255).astype(np.uint8))158 159    return mask_display, mask_overlay_display, combined_display, confidences_dict, confidence, pred_name160 161 162# --- STREAMLIT UI ---163st.title("AI Wound Classifier")164st.markdown("Upload a wound image and the AI model will automatically segment and classify the wound type.")165 166col1, col2 = st.columns([1, 1])167 168with col1:169    uploaded_file = st.file_uploader("Upload Wound Image", type=["jpg", "jpeg", "png"])170    if uploaded_file is not None:171        image = Image.open(uploaded_file)172        st.image(image, caption="Uploaded Image", use_container_width=True)173        analyze_button = st.button("🔍 Analyze Image", type="primary", use_container_width=True)174 175if uploaded_file is not None and analyze_button:176    with st.spinner('Analyzing...'):177        mask_display, mask_overlay_display, combined_display, confidences_dict, confidence, pred_name = analyze_wound(image)178 179        with col1:180            st.markdown("### Image Analysis Visualizations")181            vis1, vis2, vis3 = st.columns(3)182            with vis1:183                st.image(mask_display, caption="Segmentation Mask", use_container_width=True)184            with vis2:185                st.image(mask_overlay_display, caption="Mask Overlay", use_container_width=True)186            with vis3:187                st.image(combined_display, caption="GradCAM Overlay", use_container_width=True)188 189        with col2:190            st.markdown("### Top AI Predictions")191            192            # Create a dataframe for the bar chart193            df = pd.DataFrame(list(confidences_dict.items()), columns=['Class', 'Probability'])194            df = df.sort_values(by='Probability', ascending=False)195            196            # Display progress bars for each class197            for _, row in df.iterrows():198                st.write(f"**{row['Class']}**: {row['Probability']*100:.1f}%")199                st.progress(row['Probability'])200            201            st.markdown("### Severity Indicator")202            if confidence >= 0.8:203                st.markdown("<div class='severity-high'>🟢 <b>High Confidence</b></div>", unsafe_allow_html=True)204            elif confidence >= 0.5:205                st.markdown("<div class='severity-moderate'>🟡 <b>Moderate Confidence</b></div>", unsafe_allow_html=True)206            else:207                st.markdown("<div class='severity-low'>🔴 <b>Low Confidence</b></div>", unsafe_allow_html=True)208            209            st.markdown("### AI Diagnosis Report")210            st.markdown(f"""211            <div class='report-box'>212            <h4>AI WOUND CLASSIFICATION REPORT</h4>213            <br/>214            <b>Predicted Type:</b> {friendly_names[pred_name]}<br/>215            <b>Confidence Score:</b> {confidence:.2f}<br/>216            <br/>217            <b>Models:</b><br/>218            - Segmentation: Attention U-Net (EfficientNet-B0)<br/>219            - Classification: Soft-Guided Attention Network (EfficientNet-B3)<br/>220            <br/>221            <i>Note: This AI system is for demonstration and research purposes only.</i>222            </div>223            """, unsafe_allow_html=True)224