app.py

import streamlit as st
import torch
import torch.nn as nn
import numpy as np
import pickle
from PIL import Image
import gc
import os
from torchvision.transforms import v2 as T
import pandas as pd

# ── Page Configuration ────────────────────────────────────────────────────────
st.set_page_config(
    page_title="Cancer Histopathology Classifier",
    page_icon="🔬",
    layout="wide"
)

# ── Device (forced CPU) ──────────────────────────────────────────────────────
device = torch.device("cpu")

# ── Class names (26 cancer types) ───────────────────────────────────────────
CLASS_NAMES = [
    "all_benign", "all_early", "all_pre", "all_pro",
    "brain_glioma", "brain_menin", "brain_tumor",
    "breast_benign", "breast_malignant",
    "cervix_dyk", "cervix_koc", "cervix_mep", "cervix_pab", "cervix_sfi",
    "colon_aca", "colon_bnt",
    "kidney_normal", "kidney_tumor",
    "lung_aca", "lung_bnt", "lung_scc",
    "lymph_cll", "lymph_fl", "lymph_mcl",
    "oral_normal", "oral_scc",
]

# ── Local path to fine-tuned Qwen adapter ───────────────────────────────────
QWEN_LOCAL_PATH = "qwen-cancer-finetuned"

# ── CancerCNN ────────────────────────────────────────────────────────────────
class CancerCNN(nn.Module):
    def __init__(self, num_classes=26):
        super().__init__()
        self.layers = nn.Sequential(
            nn.Conv2d(3, 16, 3, padding=1), nn.BatchNorm2d(16), nn.ReLU(),
            nn.Conv2d(16, 16, 3, padding=1), nn.BatchNorm2d(16), nn.ReLU(),
            nn.MaxPool2d(2),
            nn.Conv2d(16, 32, 3, padding=1), nn.BatchNorm2d(32), nn.ReLU(),
            nn.MaxPool2d(2),
            nn.Conv2d(32, 32, 3, padding=1), nn.BatchNorm2d(32), nn.ReLU(),
            nn.Flatten(),
            nn.Dropout(0.3),
            nn.Linear(32 * 28 * 28, num_classes),
        )

    def forward(self, x):
        return self.layers(x)

# ── Phikon classifier head ───────────────────────────────────────────────────
class PhikonHead(nn.Module):
    def __init__(self, num_classes=26):
        super().__init__()
        self.head = nn.Linear(768, num_classes)

    def forward(self, x):
        return self.head(x)

# ── Lazy model cache ─────────────────────────────────────────────────────────
if "model_cache" not in st.session_state:
    st.session_state.model_cache = {"name": None, "model": None, "extra": None}

def _evict():
    st.session_state.model_cache["name"] = None
    st.session_state.model_cache["model"] = None
    st.session_state.model_cache["extra"] = None
    gc.collect()

# ── Transforms ───────────────────────────────────────────────────────────────
cnn_transform = T.Compose([
    T.Resize((112, 112)),
    T.ToImage(),
    T.ToDtype(torch.float32, scale=True),
    T.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5]),
])

phikon_transform = T.Compose([
    T.Resize((224, 224)),
    T.ToImage(),
    T.ToDtype(torch.float32, scale=True),
    T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
])

# ── Model loaders ─────────────────────────────────────────────────────────────
def load_cnn():
    if st.session_state.model_cache["name"] == "cnn":
        return st.session_state.model_cache["model"]
    _evict()
    m = CancerCNN(num_classes=26)
    m.load_state_dict(torch.load("cancer_cnn.pt", map_location=device), strict=False)
    m.eval().to(device)
    st.session_state.model_cache["name"] = "cnn"
    st.session_state.model_cache["model"] = m
    return m

def load_svm():
    if st.session_state.model_cache["name"] == "svm":
        return st.session_state.model_cache["model"], st.session_state.model_cache["extra"]
    _evict()
    with open("svm_model.pkl", "rb") as f:
        svm = pickle.load(f)
    from img2vec_pytorch import Img2Vec
    img2vec = Img2Vec(cuda=False)
    st.session_state.model_cache["name"] = "svm"
    st.session_state.model_cache["model"] = svm
    st.session_state.model_cache["extra"] = img2vec
    return svm, img2vec

def load_qwen():
    if st.session_state.model_cache["name"] == "qwen":
        return st.session_state.model_cache["model"], st.session_state.model_cache["extra"]
    _evict()

    if not os.path.isdir(QWEN_LOCAL_PATH):
        raise FileNotFoundError(
            f"Expected local Qwen adapter folder at '{QWEN_LOCAL_PATH}' but it "
            f"was not found."
        )

    from transformers import Qwen2VLForConditionalGeneration, AutoProcessor
    from peft import PeftModel

    # Load base model
    base_model_id = "Qwen/Qwen2-VL-2B-Instruct"
    base_model = Qwen2VLForConditionalGeneration.from_pretrained(
        base_model_id,
        torch_dtype=torch.float32,
        device_map=None,
    )
    
    # Layer the adapter on top
    model = PeftModel.from_pretrained(base_model, QWEN_LOCAL_PATH)
    model.to(device)
    model.eval()

    processor = AutoProcessor.from_pretrained(QWEN_LOCAL_PATH)

    st.session_state.model_cache["name"] = "qwen"
    st.session_state.model_cache["model"] = model
    st.session_state.model_cache["extra"] = processor
    return model, processor

def load_phikon():
    if st.session_state.model_cache["name"] == "phikon":
        return st.session_state.model_cache["model"], st.session_state.model_cache["extra"]
    _evict()
    from transformers import AutoModel
    backbone = AutoModel.from_pretrained("owkin/phikon").to(device)
    backbone.eval()
    head = PhikonHead(num_classes=26).to(device)
    
    ckpt = torch.load("phikon_head.pt", map_location=device)
    state_dict = ckpt.get("head_state_dict", ckpt)
    
    mapped_state_dict = {}
    for k, v in state_dict.items():
        if k == "weight":
            mapped_state_dict["head.weight"] = v
        elif k == "bias":
            mapped_state_dict["head.bias"] = v
        else:
            mapped_state_dict[k] = v
            
    head.load_state_dict(mapped_state_dict, strict=False)
    head.eval()
    
    st.session_state.model_cache["name"] = "phikon"
    st.session_state.model_cache["model"] = backbone
    st.session_state.model_cache["extra"] = head
    return backbone, head

# ── Inference functions ───────────────────────────────────────────────────────
def predict_cnn(image: Image.Image):
    model = load_cnn()
    x = cnn_transform(image).unsqueeze(0).to(device)
    with torch.no_grad():
        logits = model(x)
        probs = torch.softmax(logits, dim=1)[0]
    top3 = probs.topk(3)
    result = {CLASS_NAMES[i]: float(p) for i, p in zip(top3.indices, top3.values)}
    pred = CLASS_NAMES[probs.argmax().item()]
    return pred, result

def predict_svm(image: Image.Image):
    svm, img2vec = load_svm()
    vec = img2vec.get_vec(image, tensor=False).reshape(1, -1)
    
    raw_pred = svm.predict(vec)[0]
    try:
        pred = CLASS_NAMES[int(raw_pred)]
    except (ValueError, TypeError):
        pred = str(raw_pred)
        
    proba = svm.predict_proba(vec)[0] if hasattr(svm, "predict_proba") else None
    
    if proba is not None:
        top3_idx = np.argsort(proba)[-3:][::-1]
        result = {CLASS_NAMES[i]: float(proba[i]) for i in top3_idx}
    else:
        result = {pred: 1.0}
        
    return pred, result

def predict_phikon(image: Image.Image):
    backbone, head = load_phikon()
    x = phikon_transform(image).unsqueeze(0).to(device)
    with torch.no_grad():
        features = backbone(x).last_hidden_state[:, 0, :]
        logits = head(features)
        probs = torch.softmax(logits, dim=1)[0]
    top3 = probs.topk(3)
    result = {CLASS_NAMES[i]: float(p) for i, p in zip(top3.indices, top3.values)}
    pred = CLASS_NAMES[probs.argmax().item()]
    return pred, result

def predict_qwen_chat(chat_history):
    """Processes the full conversation history for Qwen2-VL"""
    model, processor = load_qwen()
    
    qwen_msgs = []
    images = []
    
    # 1. Figure out if the user is asking a follow-up or classifying a new image
    is_follow_up_chat = True
    if len(chat_history) > 0:
        latest_msg = chat_history[-1]
        if "image" in latest_msg:
            is_follow_up_chat = False # It has an image, so it's a classification task
    
    # 2. Rebuild the history
    for msg in chat_history:
        content = []
        if msg.get("image"):
            content.append({"type": "image", "image": msg["image"]})
            images.append(msg["image"])
            
        if msg.get("text"):
            content.append({"type": "text", "text": msg["text"]})
            
        qwen_msgs.append({"role": msg["role"], "content": content})
        
    input_text = processor.apply_chat_template(qwen_msgs, add_generation_prompt=True)
    
    if len(images) > 0:
        inputs = processor(images=images, text=input_text, return_tensors="pt").to(device)
    else:
        inputs = processor(text=input_text, return_tensors="pt").to(device)
        
    with torch.no_grad():
        # 3. THE MAGIC TRICK: Turn off the adapter for regular text chat!
        if is_follow_up_chat and hasattr(model, "disable_adapter"):
            with model.disable_adapter():
                output = model.generate(**inputs, max_new_tokens=300)
        else:
            # Leave adapter on for image classification
            output = model.generate(**inputs, max_new_tokens=200)
            
    response = processor.decode(output[0], skip_special_tokens=True)
    
    if "assistant" in response.lower():
        response = response.split("assistant")[-1].strip()
        
    return response

# ── Main standard inference dispatcher ────────────────────────────────────────
def classify(pil_image, model_choice):
    if model_choice == "CancerCNN":
        pred, probs = predict_cnn(pil_image)
        explanation = f"**CancerCNN predicted:** {pred}\n\nThis is a custom CNN trained from scratch on 26 cancer types with BatchNorm, Dropout, and data augmentation. It achieved 84.79% test accuracy."
        return pred, probs, explanation

    elif model_choice == "SVM (img2vec)":
        pred, probs = predict_svm(pil_image)
        explanation = f"**SVM predicted:** {pred}\n\nThis Support Vector Machine uses ResNet18 embeddings (via img2vec) as features. Classical ML approach — no deep learning training required."
        return pred, probs, explanation

    elif model_choice == "Phikon (ViT-B Histopathology)":
        pred, probs = predict_phikon(pil_image)
        explanation = f"**Phikon predicted:** {pred}\n\nPhikon is a ViT-Base model pretrained on 40M pan-cancer histopathology tiles from TCGA using self-supervised learning. It extracts domain-specific features far beyond what ImageNet-pretrained models can capture."
        return pred, probs, explanation

    return "Unknown model", {}, ""


# ── Streamlit UI ──────────────────────────────────────────────────────────────
st.title("🔬 Cancer Histopathology Classifier")
st.markdown("""
Upload a histopathology image and select a model to classify the cancer type across 26 categories.
⚠️ *Running on CPU — Deep learning inference may take a moment.*
""")

# High-level model selection controls the entire UI layout
model_choice = st.selectbox(
    "Select Model",
    ["CancerCNN", "SVM (img2vec)", "Phikon (ViT-B Histopathology)", "Qwen2-VL-2B (Fine-tuned)"]
)
st.markdown("---")

if model_choice == "Qwen2-VL-2B (Fine-tuned)":
    # ── QWEN CHAT UI (GEMINI STYLE) ───────────────────────────────────────────
    st.subheader("💬 Qwen2-VL Analysis Chat")
    
    # Initialize chat history
    if "qwen_messages" not in st.session_state:
        st.session_state.qwen_messages = []
        
    # Display existing chat history
    chat_container = st.container()
    with chat_container:
        for msg in st.session_state.qwen_messages:
            with st.chat_message(msg["role"]):
                if msg.get("image"):
                    st.image(msg["image"], width=300)
                if msg.get("text"):
                    st.markdown(msg["text"])
                    
    # The Attachment & Chat Input Area
    st.caption("📎 Attach an image for your next message:")
    chat_image_upload = st.file_uploader("Upload Image", type=["jpg", "jpeg", "png"], label_visibility="collapsed")
    
    # Input box (with the auto-prompt injection logic we discussed)
    if prompt := st.chat_input("Message Qwen... (e.g., 'What type of cancer is this?')"):
        
        # If they just uploaded an image and hit enter without typing, auto-fill the training prompt
        if not prompt.strip() and chat_image_upload is not None:
            prompt = "What type of cancer is shown in this histopathology image?"
            
        user_msg = {"role": "user", "text": prompt}
        
        # Attach image only if it's new (prevents re-uploading the same image on follow-up questions)
        if chat_image_upload is not None:
            file_sig = f"{chat_image_upload.name}_{chat_image_upload.size}"
            if st.session_state.get("last_uploaded_qwen_file") != file_sig:
                user_msg["image"] = Image.open(chat_image_upload).convert("RGB")
                st.session_state.last_uploaded_qwen_file = file_sig
                
        # Append and display user message
        st.session_state.qwen_messages.append(user_msg)
        with chat_container:
            with st.chat_message("user"):
                if user_msg.get("image"):
                    st.image(user_msg["image"], width=300)
                st.markdown(prompt)
                
            # Generate and display assistant response
            with st.chat_message("assistant"):
                with st.spinner("Qwen is analyzing..."):
                    try:
                        response = predict_qwen_chat(st.session_state.qwen_messages)
                        st.markdown(response)
                        st.session_state.qwen_messages.append({"role": "assistant", "text": response})
                    except Exception as e:
                        st.error(f"An error occurred: {e}")

else:
    # ── STANDARD UI (CNN, SVM, PHIKON) ────────────────────────────────────────
    col1, col2 = st.columns(2)

    with col1:
        st.subheader("Input")
        uploaded_file = st.file_uploader("Upload Histopathology Image", type=["jpg", "jpeg", "png"])
        
        classify_btn = st.button("Classify Image", type="primary", use_container_width=True)
        
        if uploaded_file is not None:
            image = Image.open(uploaded_file).convert("RGB")
            st.image(image, caption="Uploaded Image", use_column_width=True)

    with col2:
        st.subheader("Results")
        if classify_btn:
            if uploaded_file is None:
                st.warning("Please upload an image first.")
            else:
                with st.spinner(f"Running inference with {model_choice}..."):
                    try:
                        pred, probs, explanation = classify(image, model_choice)
                        
                        st.success(f"**Predicted Cancer Type:** {pred}")
                        
                        if probs:
                            st.markdown("**Top Predictions:**")
                            df_probs = pd.DataFrame(
                                list(probs.values()), 
                                index=list(probs.keys()), 
                                columns=["Confidence"]
                            )
                            st.bar_chart(df_probs)
                        
                        st.markdown("### Model Explanation")
                        st.info(explanation)
                        
                    except Exception as e:
                        st.error(f"An error occurred during inference: {e}")

    st.markdown("---")
    st.markdown("**Dataset:** [Multi-Cancer Dataset](https://www.kaggle.com/datasets/obulisainaren/multi-cancer) — 130K images, 26 cancer types")