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")