src/app.py

import logging
import torch
import numpy as np
import pandas as pd
import cv2
from PIL import Image
import gradio as gr

from model import MetadataMelanomaModel
from config import Config
from dataset import Transform
from file_io_manager import FileIOManager
from typing import Any, cast
from pytorch_grad_cam import EigenCAM
from pytorch_grad_cam.utils.model_targets import ClassifierOutputTarget
from pytorch_grad_cam.utils.image import show_cam_on_image

logger = logging.getLogger(__name__)


class App:
    _OOD_THRESHOLD_FALLBACK = 2000.0  # fallback if saved threshold missing

    def __init__(self):
        self._device        = Config.get_training_config()['device']
        self._transform     = Transform(train=False)
        self._io            : FileIOManager | None            = None
        self._preprocessor  : Any                            = None
        self._model         : MetadataMelanomaModel | None   = None
        self._target_layers : dict                           = {}
        self._ood_mean      : torch.Tensor | None            = None
        self._ood_cov_inv   : torch.Tensor | None            = None
        self._ood_threshold : float                          = self._OOD_THRESHOLD_FALLBACK

    def load_run(self, run_name: str) -> str:
        """Load a specific run (e.g. 'efficientnet_b0', 'experiment1')."""
        try:
            # Set experiment config so the model builds with correct architecture
            import re
            m = re.match(r'^experiment(\d+)$', run_name)
            if m:
                Config.set_experiment(int(m.group(1)))
            else:
                Config.set_experiment(None)

            io           = FileIOManager.for_run(run_name)
            preprocessor = io.load_preprocessor()
            model        = MetadataMelanomaModel.build(
                num_metadata_features=preprocessor.num_output_features
            )
            io.load_gradcam_checkpoint(model, map_location=self._device)
            model.to(self._device).eval()

            self._io            = io
            self._preprocessor  = preprocessor
            self._model         = model
            self._target_layers = self._get_target_layers()
            self._load_ood_stats(io)
            logger.info("Loaded run: %s", run_name)
            return f"Loaded: {run_name}"
        except Exception as e:
            logger.exception("Error loading run %s", run_name)
            return f"Error loading {run_name}: {e}"

    def _get_target_layers(self) -> dict:
        assert self._model is not None
        bb = self._model.image_backbone
        layers: dict = {}
        if hasattr(bb, 'conv_head'):
            layers['conv_head'] = bb.conv_head
        if hasattr(bb, 'blocks'):
            n = len(bb.blocks)
            for label, idx in [('blocks[-3].conv_pwl', n - 3), (f'blocks[{n // 2}].conv_pwl', n // 2)]:
                if idx >= 0 and hasattr(bb.blocks[idx], 'conv_pwl'):
                    layers[label] = bb.blocks[idx].conv_pwl
        if not layers:
            for name, mod in reversed(list(bb.named_modules())):
                if isinstance(mod, torch.nn.Conv2d):
                    layers[name] = mod
                    break
        logger.info("CAM layers: %s", list(layers))
        return layers

    def _load_ood_stats(self, io: FileIOManager) -> None:
        try:
            stats = io.load_ood_stats(map_location=self._device)
            self._ood_mean    = stats['mean'].to(self._device)
            self._ood_cov_inv = stats['cov_inv'].to(self._device)
            self._ood_threshold = float(stats.get('threshold', self._OOD_THRESHOLD_FALLBACK))
            logger.info("OOD stats loaded (threshold=%.1f).", self._ood_threshold)
        except Exception as e:
            logger.warning("OOD stats not available — OOD detection disabled: %s", e)
            self._ood_mean = None
            self._ood_cov_inv = None

    def _mahalanobis_distance(self, img_tensor: torch.Tensor) -> float | None:
        """Compute Mahalanobis distance of image features from training distribution.

        Based on: Lee et al., "A Simple Unified Framework for Detecting
        Out-of-Distribution Samples and Adversarial Attacks", NeurIPS 2018.
        https://arxiv.org/abs/1807.03888
        """
        if self._ood_mean is None or self._ood_cov_inv is None or self._model is None:
            return None
        with torch.no_grad():
            feats = self._model.cnn_dropout(self._model.image_backbone(img_tensor)).squeeze(0)
        diff = feats - self._ood_mean
        return float(diff @ self._ood_cov_inv @ diff)

    def _prepare_metadata(self, age=None, sex=None, site=None) -> torch.Tensor:
        defaults = Config.get_metadata_config()['defaults']
        row = {
            'age_approx':                       float(age) if age is not None else defaults['age_approx'],
            'sex':                              sex  if sex  else defaults['sex'],
            'anatom_site_general_challenge':    site if site else defaults['anatom_site_general_challenge'],
        }
        arr = self._preprocessor.transform(pd.DataFrame([row]))
        return torch.tensor(arr, dtype=torch.float32).to(self._device)

    def predict_and_visualize(self, img, age=None, sex='male', site='torso'):
        if img is None:
            return "Please upload an image", None

        if self._model is None:
            return "No model loaded. Select a run from the dropdown.", None

        try:
            img_rgb, img_resized, img_tensor = self._preprocess_image(img)

            if self._model._image_only:
                metadata_tensor = torch.zeros(1, self._model.num_metadata_features,
                                              device=self._device)
            else:
                metadata_tensor = self._prepare_metadata(age, sex, site)

            ood_dist = self._mahalanobis_distance(img_tensor)
            if ood_dist is not None and ood_dist > self._ood_threshold:
                ood_warning = (f"⚠️ This image appears to be out-of-distribution "
                               f"(distance: {ood_dist:.0f}).\n"
                               f"The model is trained on skin lesion images only.\n\n")
            else:
                ood_warning = ""

            prediction_text = self._run_tta(img_rgb, metadata_tensor)

            if prediction_text is None:
                return "Error: TTA failed for all augmentations.", None

            default_layer = self._target_layers.get('blocks[-3].conv_pwl') and 'blocks[-3].conv_pwl' or next(iter(self._target_layers))
            _, side_by_side = self._run_eigencam(img_resized, img_tensor, default_layer)
            return ood_warning + prediction_text, side_by_side

        except Exception:
            logger.exception("Prediction failed")
            return "Error", None

    def _preprocess_image(self, img) -> tuple:
        img_rgb     = img.convert('RGB')
        size        = Config.get_model_config()['image_size']
        img_resized = img_rgb.resize((size, size), Image.Resampling.LANCZOS)
        img_tensor  = self._transform(img_resized).unsqueeze(0).to(self._device)
        return img_rgb, img_resized, img_tensor

    def _run_tta(self, img_rgb, metadata_tensor) -> str | None:
        assert self._model is not None
        probs: list[float] = []
        logger.info("Starting TTA with %d augmentations...", len(Transform.tta_transforms))
        for aug_name, tta_fn in Transform.tta_transforms.items():
            try:
                aug_tensor = self._transform(tta_fn(img_rgb)).unsqueeze(0).to(self._device)
                with torch.no_grad():
                    prob = torch.sigmoid(self._model(aug_tensor, metadata_tensor)).item()
                probs.append(prob)
                logger.debug("  TTA - %s: %.4f", aug_name, prob)
            except Exception as e:
                logger.warning("  TTA %s failed: %s", aug_name, e)

        if not probs:
            return None

        final_prob = float(np.mean(probs))
        logger.info("TTA averaged probability: %.4f over %d augmentations", final_prob, len(probs))
        benign = (1 - final_prob) * 100
        malignant = final_prob * 100
        return f"Benign:     {benign:.1f}%\nMalignant:  {malignant:.1f}%"

    def _run_eigencam(self, img_resized, img_tensor, target_layer_name: str) -> tuple:
        assert self._model is not None
        if target_layer_name not in self._target_layers:
            target_layer_name = next(iter(self._target_layers))
            logger.warning("Layer not found; falling back to '%s'", target_layer_name)
        try:
            cam       = EigenCAM(model=self._model.image_backbone, target_layers=[self._target_layers[target_layer_name]])
            grayscale = cam(input_tensor=img_tensor, targets=cast(Any, [ClassifierOutputTarget(0)]))[0]
            lo, hi    = grayscale.min(), grayscale.max()
            grayscale = (grayscale - lo) / (hi - lo) if hi != lo else grayscale
            rgb_f     = np.array(img_resized).astype(np.float32) / 255.0
            cam_img   = (show_cam_on_image(rgb_f, grayscale, use_rgb=True, colormap=cv2.COLORMAP_JET, image_weight=0.5) * 255).astype(np.uint8)
            return cam_img, np.hstack((np.array(img_resized), cam_img))
        except Exception:
            logger.exception("EigenCAM failed")
            return None, None


    def build_interface(self) -> gr.Blocks:
        available_runs = FileIOManager.list_available_runs()
        default_run    = Config.MODEL_NAME if Config.MODEL_NAME in available_runs else (available_runs[0] if available_runs else None)

        with gr.Blocks(title="Melanoma Detection") as iface:
            gr.Markdown("## Melanoma Detection")

            # ── Model selection + status row ──────────────────────────────
            with gr.Row():
                run_dropdown = gr.Dropdown(
                    label="Run",
                    choices=available_runs,
                    value=default_run,
                    interactive=True,
                )
                model_status = gr.Textbox(
                    label="Status", interactive=False,
                )

            # ── Main content ─────────────────────────────────────────────
            with gr.Row():

                # Left: image + metadata
                with gr.Column(scale=1):
                    image_input = gr.Image(type="pil", label="Skin Lesion Image")

                    with gr.Group():
                        gr.Markdown("**Patient Metadata**")
                        age_input  = gr.Number(label="Age", value=50)
                        sex_input  = gr.Dropdown(
                            label="Sex", choices=["male", "female", "unknown"], value="male"
                        )
                        site_input = gr.Dropdown(
                            label="Anatomical Site",
                            choices=["torso", "lower extremity", "upper extremity",
                                     "head/neck", "palms/soles", "oral/genital",
                                     "anterior torso", "posterior torso",
                                     "lateral torso", "unknown"],
                            value="torso",
                        )
                    predict_btn = gr.Button("Analyse", variant="primary")

                # Right: prediction + heatmap
                with gr.Column(scale=2):
                    prediction_output = gr.Textbox(
                        label="Prediction", lines=2, interactive=False,
                    )
                    heatmap_output = gr.Image(
                        label="Original vs. Heatmap",
                    )

            # ── Events ───────────────────────────────────────────────────
            iface.load(
                fn=lambda: self.load_run(default_run) if default_run else "No trained runs found in output/",
                outputs=[model_status],
            )

            run_dropdown.change(
                fn=self.load_run,
                inputs=[run_dropdown],
                outputs=[model_status],
            )

            predict_btn.click(
                fn=self.predict_and_visualize,
                inputs=[image_input, age_input, sex_input, site_input],
                outputs=[prediction_output, heatmap_output],
            )

        return iface

    def launch(self, **kwargs) -> None:
        self.build_interface().launch(**kwargs)