rath1108/crcc

1import streamlit as st2import torch3import torch.nn as nn4from torchvision import models, transforms5from torchvision.models import DenseNet121_Weights, DenseNet169_Weights6from PIL import Image7import numpy as np8import cv29from ultralytics import YOLO10from huggingface_hub import hf_hub_download11 12# =====================13# 🔹 Load Models14# =====================15device = torch.device("cuda" if torch.cuda.is_available() else "cpu")16 17# Download models from Hugging Face Hub18yolo_path = hf_hub_download("rath1108/crcc", "best.pt")19col_path = hf_hub_download("rath1108/crcc", "best_densenet121_balanced.pth")20hist_path = hf_hub_download("rath1108/crcc", "best_densenet121_stage.pth")21 22# Colonoscopy YOLO model (polyp detection)23yolo_model = YOLO(yolo_path)24 25# Colonoscopy DenseNet (classification)26col_model = models.densenet169(weights=DenseNet169_Weights.IMAGENET1K_V1)27col_model.classifier = nn.Linear(col_model.classifier.in_features, 2)28col_model.load_state_dict(torch.load(col_path, map_location=device))29col_model = col_model.to(device).eval()30col_class_names = ["hyperplastic", "adenomatous"]31 32# Histopathology DenseNet (stage classification)33hist_model = models.densenet121(weights=DenseNet121_Weights.IMAGENET1K_V1)34hist_model.classifier = nn.Linear(hist_model.classifier.in_features, 4)  # adjust for your dataset35hist_model.load_state_dict(torch.load(hist_path, map_location=device))36hist_model = hist_model.to(device).eval()37hist_class_names = ["Not_specified", "Stage_I", "Stage_II", "Stage_IIA", "Stage_IIB", "Stage_III", "Stage_IIIA", "Stage_IIIB", "Stage_IIIC", "Stage_IV"]  38 39 40# =====================41# 🔹 Preprocessing42# =====================43eval_transform = transforms.Compose([44    transforms.Resize((224, 224)),45    transforms.ToTensor(),46    transforms.Normalize([0.485, 0.456, 0.406],47                         [0.229, 0.224, 0.225])48])49 50# =====================51# 🔹 Prediction Functions52# =====================53def predict_colonoscopy(img: Image.Image):54    results = yolo_model.predict(img, imgsz=640, conf=0.25, save=False)55    boxes = results[0].boxes.xyxy.cpu().numpy().astype(int)56 57    img_cv = np.array(img)58    img_cv = cv2.cvtColor(img_cv, cv2.COLOR_RGB2BGR)59 60    predictions = []61    for (x1, y1, x2, y2) in boxes:62        crop = img.crop((x1, y1, x2, y2))63        crop_tensor = eval_transform(crop).unsqueeze(0).to(device)64        with torch.no_grad():65            outputs = col_model(crop_tensor)66            probs = torch.softmax(outputs, dim=1)67            conf, pred_class = torch.max(probs, 1)68 69        label = col_class_names[pred_class.item()]70        confidence = conf.item() * 10071        predictions.append((label, confidence))72 73        cv2.rectangle(img_cv, (x1, y1), (x2, y2), (0, 255, 0), 2)74        cv2.putText(img_cv, f"{label} {confidence:.1f}%", (x1, y1 - 10),75                    cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 0), 2)76 77    return cv2.cvtColor(img_cv, cv2.COLOR_BGR2RGB), predictions78 79def predict_histopathology(img: Image.Image):80    img_tensor = eval_transform(img).unsqueeze(0).to(device)81    with torch.no_grad():82        outputs = hist_model(img_tensor)83        probs = torch.softmax(outputs, dim=1)84        conf, pred_class = torch.max(probs, 1)85    return hist_class_names[pred_class.item()], conf.item() * 10086 87# =====================88# 🔹 Streamlit UI89# =====================90st.title("🩺 AI for Polyp Detection & Cancer Staging")91 92page = st.sidebar.radio("Select Image Type", ["Home", "Colonoscopy", "Histopathology"])93 94if page == "Home":95    st.write("### Welcome! 👋")96    st.write("Select **Colonoscopy** or **Histopathology** from the sidebar to continue.")97 98elif page == "Colonoscopy":99    st.header("🔎 Colonoscopy Analysis (YOLO + DenseNet)")100    uploaded_file = st.file_uploader("Upload Colonoscopy Image", type=["jpg", "png", "jpeg"])101    if uploaded_file:102        img = Image.open(uploaded_file).convert("RGB")103        st.image(img, caption="Uploaded Image", use_column_width=True)104 105        if st.button("Run Analysis"):106            output_img, preds = predict_colonoscopy(img)107            st.image(output_img, caption="Detection + Classification", use_column_width=True)108            st.write("### Predictions:")109            for i, (label, conf) in enumerate(preds):110                st.write(f"Polyp {i+1}: **{label}** ({conf:.2f}%)")111 112    if st.button("⬅️ Return Home"):113        st.rerun()114 115elif page == "Histopathology":116    st.header("🧬 Histopathology Stage Classification (DenseNet121)")117    uploaded_file = st.file_uploader("Upload Histopathology Image", type=["jpg", "png", "jpeg"])118    if uploaded_file:119        img = Image.open(uploaded_file).convert("RGB")120        st.image(img, caption="Uploaded Image", use_column_width=True)121 122        if st.button("Run Stage Prediction"):123            label, conf = predict_histopathology(img)124            st.success(f"Predicted Stage: **{label}** ({conf:.2f}%)")125 126    if st.button("⬅️ Return Home"):127        st.rerun()128