Upload scripts/predict.py with huggingface_hub

scripts/predict.py

@@ -0,0 +1,147 @@
+"""DiaFoot.AI v2 — Inference on New Images.
+
+Run the trained pipeline on any foot image.
+
+Usage:
+    python scripts/predict.py --image path/to/foot_image.jpg
+    python scripts/predict.py --image path/to/image.jpg --save-mask output_mask.png
+"""
+
+from __future__ import annotations
+
+import argparse
+import logging
+import sys
+from pathlib import Path
+
+import cv2
+import numpy as np
+import torch
+
+sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
+
+from src.data.augmentation import get_val_transforms
+from src.models.classifier import TriageClassifier
+from src.models.unetpp import build_unetpp
+
+CLASS_NAMES = {0: "Healthy", 1: "Non-DFU Wound", 2: "DFU (Diabetic Foot Ulcer)"}
+
+
+def load_and_preprocess(image_path: str) -> tuple[np.ndarray, torch.Tensor]:
+    """Load image and prepare for inference."""
+    image = cv2.imread(image_path)
+    if image is None:
+        msg = f"Cannot read image: {image_path}"
+        raise FileNotFoundError(msg)
+
+    # Keep original for display
+    original = image.copy()
+
+    # Preprocess for model
+    image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+    image_resized = cv2.resize(image_rgb, (512, 512))
+
+    transform = get_val_transforms()
+    transformed = transform(image=image_resized)
+    tensor = transformed["image"].unsqueeze(0)
+
+    return original, tensor
+
+
+def main() -> None:
+    """Run inference on a single image."""
+    parser = argparse.ArgumentParser(description="DiaFoot.AI v2 — Predict")
+    parser.add_argument("--image", type=str, required=True, help="Path to foot image")
+    parser.add_argument(
+        "--classifier-checkpoint",
+        type=str,
+        default="checkpoints/classifier/best_epoch004_1.0000.pt",
+    )
+    parser.add_argument(
+        "--segmenter-checkpoint",
+        type=str,
+        default="checkpoints/segmentation/best_epoch019_0.6781.pt",
+    )
+    parser.add_argument("--save-mask", type=str, default=None, help="Save segmentation mask")
+    parser.add_argument("--device", type=str, default="cpu")
+    args = parser.parse_args()
+
+    logging.basicConfig(level=logging.INFO, format="%(message)s")
+
+
+    device = torch.device(
+        args.device if torch.cuda.is_available() or args.device == "cpu" else "cpu"
+    )
+
+    # Load image
+    original, tensor = load_and_preprocess(args.image)
+    tensor = tensor.to(device)
+
+    print(f"\n{'=' * 50}")
+    print("DiaFoot.AI v2 — Inference")
+    print(f"Image: {args.image}")
+    print(f"{'=' * 50}")
+
+    # Step 1: Classification
+    classifier_path = Path(args.classifier_checkpoint)
+    if classifier_path.exists():
+        classifier = TriageClassifier(
+            backbone="tf_efficientnetv2_m", num_classes=3, pretrained=False
+        )
+        ckpt = torch.load(str(classifier_path), map_location="cpu", weights_only=True)
+        classifier.load_state_dict(ckpt["model_state_dict"])
+        classifier = classifier.to(device).eval()
+
+        with torch.no_grad():
+            logits = classifier(tensor)
+            probs = torch.softmax(logits, dim=1).squeeze().cpu().numpy()
+
+        pred_class = int(probs.argmax())
+        confidence = float(probs.max())
+
+        print(f"\n  Classification: {CLASS_NAMES[pred_class]}")
+        print(f"  Confidence: {confidence:.1%}")
+        for i, name in CLASS_NAMES.items():
+            print(f"    {name}: {probs[i]:.1%}")
+    else:
+        print(f"\n  Classifier checkpoint not found: {classifier_path}")
+        pred_class = 2  # Assume DFU for segmentation
+
+    # Step 2: Segmentation (if wound detected)
+    segmenter_path = Path(args.segmenter_checkpoint)
+    if pred_class in (1, 2) and segmenter_path.exists():
+        segmenter = build_unetpp(encoder_name="efficientnet-b4", encoder_weights=None, classes=1)
+        ckpt = torch.load(str(segmenter_path), map_location="cpu", weights_only=True)
+        segmenter.load_state_dict(ckpt["model_state_dict"])
+        segmenter = segmenter.to(device).eval()
+
+        with torch.no_grad():
+            seg_logits = segmenter(tensor)
+            seg_prob = torch.sigmoid(seg_logits).squeeze().cpu().numpy()
+            seg_mask = (seg_prob > 0.5).astype(np.uint8)
+
+        wound_pixels = seg_mask.sum()
+        total_pixels = seg_mask.shape[0] * seg_mask.shape[1]
+        coverage = wound_pixels / total_pixels * 100
+        area_mm2 = wound_pixels * 0.5 * 0.5  # Assuming 0.5mm/pixel
+
+        print("\n  Segmentation:")
+        print(f"    Wound detected: {'Yes' if wound_pixels > 0 else 'No'}")
+        print(f"    Wound pixels: {wound_pixels:,}")
+        print(f"    Coverage: {coverage:.1f}%")
+        print(f"    Estimated area: {area_mm2:.1f} mm2")
+
+        if args.save_mask:
+            mask_resized = cv2.resize(seg_mask * 255, (original.shape[1], original.shape[0]))
+            cv2.imwrite(args.save_mask, mask_resized)
+            print(f"    Mask saved to: {args.save_mask}")
+    elif pred_class == 0:
+        print("\n  Segmentation: Skipped (healthy foot detected)")
+    else:
+        print(f"\n  Segmenter checkpoint not found: {segmenter_path}")
+
+    print(f"\n{'=' * 50}\n")
+
+
+if __name__ == "__main__":
+    main()