diff --git a/fastapi_app.py b/fastapi_app.py index c80f551..9d5d71c 100644 --- a/fastapi_app.py +++ b/fastapi_app.py @@ -1033,6 +1033,56 @@ async def run_image_classification_inference( except Exception as e: return JSONResponse(status_code=500, content={"error": str(e)}) +@app.post("/explain/image-classification", summary="Explain an Image Classification Prediction (Grad-CAM)") +async def explain_image_classification( + model_checkpoint: str = Form(..., description="Path to the model checkpoint"), + image: UploadFile = File(..., description="The image to classify and explain.") +): + """ + Runs Grad-CAM on an uploaded image and returns the predicted class, its + confidence, and a URL to a heatmap overlay showing which regions of the + image drove the prediction. + """ + if not job_store.is_valid_checkpoint_path(model_checkpoint): + raise HTTPException(status_code=400, detail="Invalid or malicious model path") + + input_path = os.path.join(UPLOAD_DIR, f"{uuid.uuid4()}_{image.filename}") + try: + with open(input_path, "wb") as buffer: + shutil.copyfileobj(image.file, buffer) + + output_filename = f"cam_{os.path.basename(input_path)}.png" + output_path = os.path.join(OUTPUT_DIR, output_filename) + + cmd = [ + sys.executable, "image_classification_explain.py", + "--model_checkpoint", model_checkpoint, + "--image_path", input_path, + "--output_path", output_path, + ] + process = subprocess.run(cmd, capture_output=True, text=True, check=True) + + # The script emits a "RESULT_JSON:{...}" line with the prediction. + result = {} + for line in process.stdout.splitlines(): + if line.startswith("RESULT_JSON:"): + result = json.loads(line[len("RESULT_JSON:"):]) + break + + return { + "prediction": result.get("prediction"), + "confidence": result.get("confidence"), + "output_url": f"/outputs/{output_filename}", + } + except subprocess.CalledProcessError as e: + return JSONResponse(status_code=500, content={"error": "Explanation script failed", "details": e.stderr}) + except Exception as e: + return JSONResponse(status_code=500, content={"error": str(e)}) + finally: + # Clean up the uploaded input image; the overlay in OUTPUT_DIR is served. + if os.path.exists(input_path): + os.remove(input_path) + @app.post("/inference/image-segmentation", summary="Run Image Segmentation Inference") async def run_semantic_segmentation_inference( model_checkpoint: str = Form(..., description="Path to the model checkpoint"), diff --git a/image_classification_explain.py b/image_classification_explain.py new file mode 100644 index 0000000..6c6363e --- /dev/null +++ b/image_classification_explain.py @@ -0,0 +1,117 @@ +# image_classification_explain.py +# +# Grad-CAM explainability for image-classification models. Produces a heatmap +# overlay showing which regions of the image drove the predicted class, so a +# no-code user can see *why* a prediction was made — not just the label. +# +# Works across the registry's two model families: +# - CNNs (ResNet, ConvNeXt, EfficientNet): target the last Conv2d layer. +# - Transformers (ViT, BEiT, Swin, DINOv2): target the final LayerNorm and +# reshape the token sequence back to a 2D grid. + +import argparse +import json + +import numpy as np +import torch +from PIL import Image +from transformers import AutoModelForImageClassification, AutoImageProcessor + +from pytorch_grad_cam import GradCAM +from pytorch_grad_cam.utils.image import show_cam_on_image +from pytorch_grad_cam.utils.model_targets import ClassifierOutputTarget + + +class HFClassifierWrapper(torch.nn.Module): + """Return raw logits (a plain tensor) so Grad-CAM can call .backward().""" + + def __init__(self, model): + super().__init__() + self.model = model + + def forward(self, x): + return self.model(pixel_values=x).logits + + +def _grid_reshape(tensor): + """Reshape a (B, tokens, C) transformer output to (B, C, H, W). + + Drops the leading CLS token when the token count is a perfect square + 1; + infers a square grid from the number of patch tokens. + """ + n_tokens = tensor.size(1) + n_patches = n_tokens - 1 if int((n_tokens - 1) ** 0.5) ** 2 == (n_tokens - 1) else n_tokens + start = n_tokens - n_patches + side = int(round(n_patches ** 0.5)) + result = tensor[:, start:start + side * side, :].reshape( + tensor.size(0), side, side, tensor.size(2) + ) + return result.permute(0, 3, 1, 2) + + +def select_target_layer(model): + """Pick a Grad-CAM target layer and (optionally) a reshape transform. + + Returns (target_layers, reshape_transform_or_None). + """ + conv_layers = [m for m in model.modules() if isinstance(m, torch.nn.Conv2d)] + if conv_layers: + # CNN: the last conv layer holds the richest localized features. + return [conv_layers[-1]], None + + # Transformer: use the last LayerNorm and reshape tokens to a grid. + norm_layers = [m for m in model.modules() if isinstance(m, torch.nn.LayerNorm)] + if not norm_layers: + raise ValueError("Could not find a Conv2d or LayerNorm target layer for Grad-CAM.") + return [norm_layers[-1]], _grid_reshape + + +def run_explain(model_checkpoint: str, image_path: str, output_path: str): + device = "cuda" if torch.cuda.is_available() else "cpu" + print(f"Using device: {device}") + + processor = AutoImageProcessor.from_pretrained(model_checkpoint) + model = AutoModelForImageClassification.from_pretrained(model_checkpoint).to(device).eval() + + image = Image.open(image_path).convert("RGB") + inputs = processor(images=image, return_tensors="pt").to(device) + pixel_values = inputs["pixel_values"] + + # Predicted class. + with torch.no_grad(): + logits = model(pixel_values=pixel_values).logits + pred_idx = int(logits.argmax(-1).item()) + pred_label = model.config.id2label.get(pred_idx, str(pred_idx)) + confidence = float(torch.softmax(logits, dim=-1)[0, pred_idx].item()) + print(f"Predicted class: {pred_label} (Confidence: {confidence:.4f})") + + # Grad-CAM. + target_layers, reshape = select_target_layer(model) + wrapper = HFClassifierWrapper(model) + cam = GradCAM(model=wrapper, target_layers=target_layers, reshape_transform=reshape) + grayscale_cam = cam( + input_tensor=pixel_values, + targets=[ClassifierOutputTarget(pred_idx)], + )[0] + + # Overlay the heatmap on the (normalized-to-0..1) processed image so the CAM + # aligns with what the model actually saw. + vis = pixel_values[0].detach().cpu().numpy().transpose(1, 2, 0) + vis = (vis - vis.min()) / (vis.max() - vis.min() + 1e-8) + overlay = show_cam_on_image(vis, grayscale_cam, use_rgb=True) + Image.fromarray(overlay).save(output_path) + print(f"Explanation saved to: {output_path}") + + return {"prediction": pred_label, "confidence": confidence, "output_path": output_path} + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="Grad-CAM explanation for image classification.") + parser.add_argument("--model_checkpoint", type=str, required=True) + parser.add_argument("--image_path", type=str, required=True) + parser.add_argument("--output_path", type=str, required=True) + args = parser.parse_args() + + result = run_explain(args.model_checkpoint, args.image_path, args.output_path) + # Emit a JSON line so the API can parse the prediction alongside the image. + print("RESULT_JSON:" + json.dumps(result)) diff --git a/requirements.txt b/requirements.txt index 7b7ee79..d982740 100644 --- a/requirements.txt +++ b/requirements.txt @@ -30,6 +30,7 @@ pycocotools==2.0.11 ultralytics==8.4.92 faster-coco-eval==1.7.2 segmentation-models-pytorch==0.5.0 +grad-cam==1.5.5 # --- API --- fastapi==0.139.0 diff --git a/scripts/smoke_test.py b/scripts/smoke_test.py index 163869d..65da2de 100644 --- a/scripts/smoke_test.py +++ b/scripts/smoke_test.py @@ -114,7 +114,13 @@ def smoke_classification(work: Path, env): img = next((data / "test").rglob("*.png")) ok = run([PY, "image_classification_inference.py", "--model_checkpoint", ckpt, "--image_path", img], env) - return ("classification", ok and ckpt is not None) + if ok and ckpt: + # Grad-CAM explanation on the same checkpoint. + img = next((data / "test").rglob("*.png")) + ok = run([PY, "image_classification_explain.py", "--model_checkpoint", ckpt, + "--image_path", img, "--output_path", work / "cls_cam.png"], env) + ok = ok and (work / "cls_cam.png").exists() + return ("classification+explain", ok and ckpt is not None) def smoke_detection_hf(work: Path, env):