import os import io import re import base64 import numpy as np from PIL import Image from fastapi import FastAPI, File, UploadFile, HTTPException from fastapi.middleware.cors import CORSMiddleware import uvicorn import cv2 from scipy.ndimage import gaussian_filter import tensorflow as tf from tensorflow.keras.models import load_model from tensorflow.keras.utils import register_keras_serializable from transformers import TFViTModel, ViTConfig # ────────────────────────────────────────────────────────────── # 1. VIT BASE MODEL # ────────────────────────────────────────────────────────────── VIT_MODEL_NAME = "google/vit-base-patch16-224-in21k" print("Loading ViT base weights...") try: vit_base_for_loading = TFViTModel.from_pretrained(VIT_MODEL_NAME, from_pt=False) print("OK: Native TF ViT weights loaded") except Exception: try: vit_base_for_loading = TFViTModel.from_pretrained(VIT_MODEL_NAME, from_pt=True) print("OK: PyTorch -> TF ViT weights loaded") except Exception: config = ViTConfig.from_pretrained(VIT_MODEL_NAME) vit_base_for_loading = TFViTModel(config) dummy = tf.zeros((1, 3, 224, 224)) _ = vit_base_for_loading(pixel_values=dummy, training=False) print("WARN: ViT created with random weights") def vit_forward(x): x = tf.transpose(x, [0, 3, 1, 2]) outputs = vit_base_for_loading(pixel_values=x, training=False) return outputs.last_hidden_state[:, 0, :] # ────────────────────────────────────────────────────────────── # 2. CUSTOM KERAS LAYERS # ────────────────────────────────────────────────────────────── @register_keras_serializable(package='Custom', name='ViTPreprocessLayer') class ViTPreprocessLayer(tf.keras.layers.Layer): def __init__(self, **kwargs): super().__init__(**kwargs) self.mean = tf.constant([0.485, 0.456, 0.406], dtype=tf.float32, shape=[1,1,1,3]) self.std = tf.constant([0.229, 0.224, 0.225], dtype=tf.float32, shape=[1,1,1,3]) def call(self, images): return (images - self.mean) / self.std def get_config(self): return super().get_config() @register_keras_serializable(package='Custom', name='ViTFeatureExtractor') class ViTFeatureExtractor(tf.keras.layers.Layer): def __init__(self, vit_model=None, **kwargs): super().__init__(**kwargs) self.vit_model = vit_model if vit_model is not None else vit_base_for_loading def call(self, inputs, training=False): x = tf.transpose(inputs, [0, 3, 1, 2]) out = self.vit_model(pixel_values=x, training=training) return out.last_hidden_state[:, 0, :] def get_config(self): return super().get_config() CUSTOM_OBJECTS = { 'ViTPreprocessLayer': ViTPreprocessLayer, 'ViTFeatureExtractor': ViTFeatureExtractor, 'vit_forward': vit_forward, } # ────────────────────────────────────────────────────────────── # 3. LOAD MODELS (.h5 format) # ────────────────────────────────────────────────────────────── MODEL_DIR = os.path.join(os.path.dirname(__file__), "models") GAN_PATH = os.path.join(MODEL_DIR, "generator_best.h5") CLF_PATH = os.path.join(MODEL_DIR, "classifier_vit_best_joint.h5") print("Loading GAN generator...") generator = load_model(GAN_PATH, custom_objects=CUSTOM_OBJECTS, compile=False) print("OK: Generator loaded") print("Loading ViT classifier...") classifier = load_model(CLF_PATH, custom_objects=CUSTOM_OBJECTS, compile=False) print("OK: Classifier loaded") # ────────────────────────────────────────────────────────────── # 4. CLASS NAMES # ────────────────────────────────────────────────────────────── CLASS_NAMES_RAW = [ "1Gallstones", "2Abdomen and retroperitoneum", "3cholecystitis", "4Membranous and gangrenous cholecystitis", "5Perforation", "6Polyps and cholesterol crystals", "7Adenomyomatosis", "8Carcinoma", "9Various causes of gallbladder wall thickening", ] def clean_class_name(raw: str) -> str: """Remove leading numeric prefix and ensure title-cased first letter.""" cleaned = re.sub(r'^\d+', '', raw).strip() return cleaned[0].upper() + cleaned[1:] if cleaned else cleaned CLASS_NAMES_DISPLAY = [clean_class_name(c) for c in CLASS_NAMES_RAW] # ────────────────────────────────────────────────────────────── # 5. IMAGE HELPERS # ────────────────────────────────────────────────────────────── IMG_SIZE = 224 def preprocess_image(image_bytes: bytes) -> np.ndarray: img = Image.open(io.BytesIO(image_bytes)).convert("RGB") img = img.resize((IMG_SIZE, IMG_SIZE), Image.BILINEAR) return np.array(img, dtype="float32") / 255.0 def array_to_base64(arr: np.ndarray) -> str: img = Image.fromarray((np.clip(arr, 0, 1) * 255).astype(np.uint8)) buf = io.BytesIO() img.save(buf, format="JPEG", quality=85) return "data:image/jpeg;base64," + base64.b64encode(buf.getvalue()).decode() def uint8_to_base64(arr: np.ndarray) -> str: img = Image.fromarray(arr) buf = io.BytesIO() img.save(buf, format="JPEG", quality=85) return "data:image/jpeg;base64," + base64.b64encode(buf.getvalue()).decode() # ────────────────────────────────────────────────────────────── # 6. ATTENTION ROLLOUT (XAI) # ────────────────────────────────────────────────────────────── def compute_attention_rollout(denoised_arr: np.ndarray) -> str: mean = np.array([0.485, 0.456, 0.406], dtype=np.float32) std = np.array([0.229, 0.224, 0.225], dtype=np.float32) normalized = (denoised_arr[np.newaxis] - mean) / std pixel_values = tf.transpose(normalized, [0, 3, 1, 2]) outputs = vit_base_for_loading( pixel_values=pixel_values, training=False, output_attentions=True ) rollout = tf.eye(197, dtype=tf.float32) for attn in outputs.attentions: attn_avg = tf.reduce_mean(attn[0], axis=0) attn_avg = attn_avg + tf.eye(197, dtype=tf.float32) attn_avg = attn_avg / tf.reduce_sum(attn_avg, axis=-1, keepdims=True) rollout = tf.matmul(attn_avg, rollout) cls_attention = rollout[0, 1:] mask = tf.reshape(cls_attention, (14, 14)).numpy() mask = (mask - mask.min()) / (mask.max() - mask.min() + 1e-8) smooth = cv2.resize(mask, (IMG_SIZE, IMG_SIZE), interpolation=cv2.INTER_CUBIC) smooth = gaussian_filter(smooth, sigma=4) smooth = (smooth - smooth.min()) / (smooth.max() - smooth.min() + 1e-8) heatmap_color = cv2.applyColorMap(np.uint8(255 * smooth), cv2.COLORMAP_JET) heatmap_color = cv2.cvtColor(heatmap_color, cv2.COLOR_BGR2RGB) original_uint8 = np.uint8(255 * np.clip(denoised_arr, 0, 1)) overlay = cv2.addWeighted(original_uint8, 0.45, heatmap_color, 0.55, 0) return uint8_to_base64(overlay) # ────────────────────────────────────────────────────────────── # 7. FASTAPI APP # ────────────────────────────────────────────────────────────── app = FastAPI(title="Gallbladder Classifier API", version="1.0.0") app.add_middleware( CORSMiddleware, allow_origins=["*"], allow_methods=["POST", "GET"], allow_headers=["*"], ) @app.get("/") def root(): return {"status": "ok", "message": "Gallbladder Classifier API is running"} @app.get("/classes") def get_classes(): return {"classes": CLASS_NAMES_DISPLAY} @app.post("/predict") async def predict(file: UploadFile = File(...)): if not file.content_type.startswith("image/"): raise HTTPException(status_code=400, detail="File must be an image (jpg, png, etc.)") image_bytes = await file.read() arr = preprocess_image(image_bytes) input_b64 = array_to_base64(arr) batch = arr[np.newaxis, ...] denoised_arr = generator.predict(batch, verbose=0)[0] denoised_arr = np.clip(denoised_arr, 0, 1).astype("float32") denoised_b64 = array_to_base64(denoised_arr) probs = classifier.predict(denoised_arr[np.newaxis, ...], verbose=0)[0] pred_idx = int(np.argmax(probs)) pred_class = CLASS_NAMES_DISPLAY[pred_idx] confidence = float(probs[pred_idx]) all_probs = [ {"class": CLASS_NAMES_DISPLAY[i], "probability": float(probs[i])} for i in np.argsort(probs)[::-1] ] attention_b64 = compute_attention_rollout(denoised_arr) return { "prediction": pred_class, "confidence": confidence, "probabilities": all_probs, "input_image": input_b64, "denoised_image": denoised_b64, "attention_overlay": attention_b64, } if __name__ == "__main__": uvicorn.run("main:app", host="0.0.0.0", port=7860, reload=False)