predict.py

"""
High-Accuracy Custom Mathematical ML Object Detection Inference Engine
Zero YOLO / Ultralytics dependencies.
Features:
- Precision Location Point Detection (Mathematical Centroid Keypoint + Halo + Crosshairs)
- Structural Tower Section Localization (Top/Cross-Arm, Mid/Lattice, Base/Legs)
- Class-Wise Folder Organization (rust, missing_bolt, veg, bird_nest, damaged_insulator, multi_defect)
- Production-Grade JSON Inspection Reporting for Enterprise Audit Trails
"""

import os
import sys
import json
import time
import uuid
import argparse
from datetime import datetime, timezone
from pathlib import Path
from typing import List, Dict, Any, Tuple

# Ensure project root is in sys.path for direct CLI execution
ROOT_DIR = Path(__file__).resolve().parent.parent
if str(ROOT_DIR) not in sys.path:
    sys.path.insert(0, str(ROOT_DIR))

import numpy as np
import torch
from PIL import Image, ImageDraw, ImageFont

from src.model import RawCustomDetector, decode_predictions, custom_nms
from src.dataset import DEFECT_CLASSES

# UTF-8 output on Windows console
if sys.platform == "win32":
    try:
        sys.stdout.reconfigure(encoding="utf-8")
    except Exception:
        pass

CLASS_METADATA = {
    0: {
        "name": "rust",
        "color": "#FF5722",
        "severity": "HIGH",
        "action": "Perform abrasive grit blast cleaning, apply zinc-rich epoxy primer and durable polyurethane topcoat."
    },
    1: {
        "name": "missing_bolt",
        "color": "#E91E63",
        "severity": "CRITICAL",
        "action": "EMERGENCY: Immediate structural fastener replacement required. Insert Grade 8.8 structural bolt and torque to spec."
    },
    2: {
        "name": "veg",
        "color": "#4CAF50",
        "severity": "MEDIUM",
        "action": "Clear vegetation canopy from steel foundation and apply soil-active herbicide barrier around tower leg stub."
    },
    3: {
        "name": "bird_nest",
        "color": "#FFC107",
        "severity": "MEDIUM",
        "action": "Evaluate electrical clearance risk, safely relocate nest materials, and install anti-perch bird deterrent cones."
    },
    4: {
        "name": "damaged_insulator",
        "color": "#9C27B0",
        "severity": "CRITICAL",
        "action": "EMERGENCY: High electrical flashover hazard. De-energize bay and perform hot-line replacement of cracked insulator discs."
    }
}


def get_tower_section(norm_x: float, norm_y: float) -> Tuple[str, str, str]:
    """
    Determines structural section of the lattice/transmission tower
    based on normalized vertical elevation (norm_y) and horizontal bay (norm_x).
    """
    if norm_y < 0.35:
        vertical_zone = "Top Section (Cross-Arm / Peak)"
        short_vert = "Top-Crossarm"
    elif norm_y <= 0.68:
        vertical_zone = "Mid Section (Lattice Body)"
        short_vert = "Mid-Lattice"
    else:
        vertical_zone = "Base Section (Legs / Foundation)"
        short_vert = "Base-Leg"

    if norm_x < 0.35:
        lateral_bay = "Left Bay"
    elif norm_x <= 0.65:
        lateral_bay = "Center Bay"
    else:
        lateral_bay = "Right Bay"

    full_section = f"{vertical_zone} - {lateral_bay}"
    short_section = f"{short_vert} | {lateral_bay}"
    return full_section, short_section, short_vert


def draw_tower_elevation_indicator(draw: ImageDraw.ImageDraw, img_w: int, img_h: int):
    """
    Renders a structural elevation guide on the left margin of the image.
    """
    bar_w = 28
    # Top Section (0% - 35%)
    draw.rectangle([0, 0, bar_w, int(img_h * 0.35)], fill="#0077BE", outline="#FFFFFF")
    # Mid Section (35% - 68%)
    draw.rectangle([0, int(img_h * 0.35), bar_w, int(img_h * 0.68)], fill="#3B5998", outline="#FFFFFF")
    # Base Section (68% - 100%)
    draw.rectangle([0, int(img_h * 0.68), bar_w, img_h], fill="#2E4053", outline="#FFFFFF")

    # Section text markers
    draw.text((3, 6), "TOP", fill="#FFFFFF")
    draw.text((3, int(img_h * 0.35) + 6), "MID", fill="#FFFFFF")
    draw.text((3, int(img_h * 0.68) + 6), "BASE", fill="#FFFFFF")


def load_custom_model(
    weights_path: str,
    device: torch.device,
    num_classes: int = 5
) -> RawCustomDetector:
    """
    Loads native PyTorch RawCustomDetector model checkpoint.
    """
    model = RawCustomDetector(num_classes=num_classes).to(device)
    if os.path.exists(weights_path):
        checkpoint = torch.load(weights_path, map_location=device)
        if isinstance(checkpoint, dict) and "model_state_dict" in checkpoint:
            model.load_state_dict(checkpoint["model_state_dict"])
            print(f"Loaded weights checkpoint from: {weights_path} (Epoch: {checkpoint.get('epoch', 'N/A')})")
        elif isinstance(checkpoint, dict):
            model.load_state_dict(checkpoint)
            print(f"Loaded raw state dict from: {weights_path}")
        else:
            print(f"Loaded weights from: {weights_path}")
    else:
        print(f"⚠️ Warning: Checkpoint not found at '{weights_path}'. Running with initialized weights.")
    model.eval()
    return model


def run_custom_prediction(
    source: str,
    weights: str = "d:/DETECTION_MODEL/runs/detect/custom_tower_defects/weights/best.pt",
    conf: float = 0.25,
    iou: float = 0.45,
    imgsz: int = 640,
    device: str = "auto",
    output_dir: str = "d:/DETECTION_MODEL/runs/detect/predictions"
) -> Dict[str, Any]:
    """
    Executes end-to-end mathematical object detection inference and reporting.
    """
    if device in ["cuda", "auto"] or str(device).startswith("cuda"):
        if torch.cuda.is_available():
            dev = torch.device(device if device != "auto" else "cuda")
        else:
            if device != "auto":
                print("⚠️ Warning: CUDA device requested, but CUDA is not available on this system. Falling back to CPU.")
            dev = torch.device("cpu")
    else:
        dev = torch.device(device)

    print("=" * 75)
    print("🔍 Running Custom Mathematical ML Object Detection Inference")
    print(f"Architecture: RawCustomDetector (Pure PyTorch & NumPy)")
    print(f"Device:       {dev}")
    print(f"Weights:      {weights}")
    print(f"Confidence:   {conf}")
    print(f"NMS IoU:      {iou}")
    print(f"Source:       {source}")
    print(f"Output Root:  {output_dir}")
    print("=" * 75)

    out_root = Path(output_dir)
    for c_id, name in DEFECT_CLASSES.items():
        (out_root / name).mkdir(parents=True, exist_ok=True)
    (out_root / "multi_defect").mkdir(parents=True, exist_ok=True)
    (out_root / "all").mkdir(parents=True, exist_ok=True)
    (out_root / "clean_no_defects").mkdir(parents=True, exist_ok=True)
    (out_root / "json").mkdir(parents=True, exist_ok=True)

    # 1. Load Model
    model = load_custom_model(weights, dev, num_classes=len(DEFECT_CLASSES))

    # 2. Gather image paths
    src_path = Path(source)
    valid_exts = {".jpg", ".jpeg", ".png", ".bmp"}
    if src_path.is_file():
        image_paths = [src_path]
    elif src_path.is_dir():
        image_paths = sorted([p for p in src_path.glob("*") if p.suffix.lower() in valid_exts])
    else:
        raise FileNotFoundError(f"Source path not found: {source}")

    if len(image_paths) == 0:
        print(f"No valid image files found in {source}")
        return {}

    total_defects = 0
    saved_destinations = {name: 0 for name in DEFECT_CLASSES.values()}
    batch_json_reports = []
    start_total_time = time.time()

    # 3. Process each image
    for img_path in image_paths:
        t0 = time.time()
        with Image.open(img_path) as raw_img:
            orig_img = raw_img.convert("RGB")
            orig_w, orig_h = orig_img.size
            file_size = img_path.stat().st_size if img_path.exists() else 0

        # Preprocessing: resize and normalize
        resized_img = orig_img.resize((imgsz, imgsz), Image.Resampling.BILINEAR)
        img_np = np.array(resized_img, dtype=np.float32) / 255.0
        img_tensor = torch.from_numpy(img_np).permute(2, 0, 1).unsqueeze(0).to(dev)

        # Mathematical Forward Pass
        with torch.no_grad():
            pred_maps = model(img_tensor)
            decoded_batch = decode_predictions(pred_maps, conf_thresh=conf)
            b_boxes, b_scores, b_classes = decoded_batch[0]
            # Mathematical Non-Maximum Suppression (NMS)
            nms_results = custom_nms(b_boxes, b_scores, b_classes, iou_thresh=iou)

        num_found = nms_results.shape[0]
        total_defects += num_found
        latency_ms = (time.time() - t0) * 1000.0
        print(f"\n📄 {img_path.name}: {num_found} defect(s) detected ({latency_ms:.1f}ms)")

        # Prepare canvas for visual rendering
        draw_img = orig_img.copy()
        draw = ImageDraw.Draw(draw_img)
        draw_tower_elevation_indicator(draw, orig_w, orig_h)

        scale_x = orig_w / float(imgsz)
        scale_y = orig_h / float(imgsz)

        detected_classes = set()
        detections_payload = []
        severity_counts = {"CRITICAL": 0, "HIGH": 0, "MEDIUM": 0, "LOW": 0}

        if num_found > 0:
            for det_idx in range(num_found):
                det = nms_results[det_idx]
                # Network coordinate space -> Original image pixel space
                x1 = float(det[0].item()) * scale_x
                y1 = float(det[1].item()) * scale_y
                x2 = float(det[2].item()) * scale_x
                y2 = float(det[3].item()) * scale_y
                real_cx = float(det[4].item()) * scale_x
                real_cy = float(det[5].item()) * scale_y
                score = float(det[6].item())
                cls_id = int(det[7].item())

                cls_name = DEFECT_CLASSES.get(cls_id, f"defect_{cls_id}")
                detected_classes.add(cls_name)

                # Clamp to image boundaries
                x1 = max(0.0, min(float(orig_w), x1))
                y1 = max(0.0, min(float(orig_h), y1))
                x2 = max(0.0, min(float(orig_w), x2))
                y2 = max(0.0, min(float(orig_h), y2))
                box_w = x2 - x1
                box_h = y2 - y1

                norm_cx = real_cx / orig_w
                norm_cy = real_cy / orig_h

                full_sec, short_sec, short_vert = get_tower_section(norm_cx, norm_cy)
                meta = CLASS_METADATA.get(cls_id, {
                    "name": cls_name,
                    "color": "#00FFFF",
                    "severity": "MEDIUM",
                    "action": "Inspect defect area and verify structural clearance."
                })
                color = meta["color"]
                severity = meta["severity"]
                severity_counts[severity] = severity_counts.get(severity, 0) + 1

                # 1. Draw Bounding Box with Corner Accents
                draw.rectangle([x1, y1, x2, y2], outline=color, width=3)
                corner_len = min(15.0, min(box_w, box_h) / 3.0)
                # Top-left corner
                draw.line([(x1, y1), (x1 + corner_len, y1)], fill=color, width=5)
                draw.line([(x1, y1), (x1, y1 + corner_len)], fill=color, width=5)
                # Bottom-right corner
                draw.line([(x2, y2), (x2 - corner_len, y2)], fill=color, width=5)
                draw.line([(x2, y2), (x2, y2 - corner_len)], fill=color, width=5)

                # 2. Draw Mathematical Centroid Marker (Halo + Crosshairs)
                pt_radius = 7
                # Halo ring
                draw.ellipse(
                    [real_cx - pt_radius, real_cy - pt_radius, real_cx + pt_radius, real_cy + pt_radius],
                    outline="#FFFFFF",
                    width=2
                )
                # Inner solid center dot
                draw.ellipse(
                    [real_cx - 2, real_cy - 2, real_cx + 2, real_cy + 2],
                    fill=color
                )
                # Precision crosshair lines
                ch_len = 12
                draw.line([(real_cx - ch_len, real_cy), (real_cx + ch_len, real_cy)], fill="#FFFFFF", width=2)
                draw.line([(real_cx, real_cy - ch_len), (real_cx, real_cy + ch_len)], fill="#FFFFFF", width=2)

                # 3. Location Point Coordinate Badge
                pt_tag = f"({int(real_cx)}, {int(real_cy)}) [{short_vert}]"
                pt_tag_w = len(pt_tag) * 8
                draw.rectangle(
                    [real_cx + 8, real_cy - 10, real_cx + 8 + pt_tag_w, real_cy + 6],
                    fill="#111111",
                    outline=color,
                    width=1
                )
                draw.text((real_cx + 12, real_cy - 9), pt_tag, fill="#FFFFFF")

                # 4. Floating Defect Tag above Bounding Box
                label_text = f"#{det_idx + 1} {cls_name.upper()} {int(score * 100)}%"
                badge_w = max(len(label_text) * 9, 130)
                badge_h = 24
                badge_y1 = max(0.0, y1 - badge_h - 2)
                badge_y2 = badge_y1 + badge_h

                draw.rectangle([x1, badge_y1, x1 + badge_w, badge_y2], fill=color)
                draw.text((x1 + 6, badge_y1 + 4), label_text, fill="#FFFFFF")

                # Record detection payload
                detections_payload.append({
                    "defect_id": f"DEF-{det_idx + 1:03d}",
                    "class_id": cls_id,
                    "class_name": cls_name,
                    "confidence": round(score, 4),
                    "confidence_pct": f"{int(score * 100)}%",
                    "severity": severity,
                    "bounding_box_pixels": {
                        "x1": round(x1, 1),
                        "y1": round(y1, 1),
                        "x2": round(x2, 1),
                        "y2": round(y2, 1),
                        "width": round(box_w, 1),
                        "height": round(box_h, 1)
                    },
                    "location_point": {
                        "centroid_x_px": int(real_cx),
                        "centroid_y_px": int(real_cy),
                        "normalized_x": round(norm_cx, 4),
                        "normalized_y": round(norm_cy, 4),
                        "crosshair_radius_px": ch_len
                    },
                    "tower_structural_section": {
                        "full_localization": full_sec,
                        "short_code": short_sec,
                        "vertical_elevation": short_vert
                    },
                    "maintenance_recommendation": meta["action"]
                })

        # 5. Header Overlay Banner
        header_h = 56
        draw.rectangle([0, 0, orig_w, header_h], fill="#1E1E24")
        draw.line([(0, header_h), (orig_w, header_h)], fill="#FF9800", width=2)

        title_text = f"TOWER DEFECT INSPECTION  |  {img_path.name}"
        draw.text((36, 10), title_text, fill="#FFFFFF")

        if num_found == 0:
            status_text = "STATUS: PASS - ZERO DEFECTS DETECTED"
            status_color = "#4CAF50"
            overall_severity = "NONE"
        else:
            defect_summary = ", ".join(sorted(list(detected_classes)))
            status_text = f"STATUS: {num_found} DEFECT(S) IDENTIFIED: [{defect_summary.upper()}]"
            if severity_counts["CRITICAL"] > 0:
                status_color = "#E91E63"
                overall_severity = "CRITICAL"
            elif severity_counts["HIGH"] > 0:
                status_color = "#FF5722"
                overall_severity = "HIGH"
            else:
                status_color = "#FFC107"
                overall_severity = "MEDIUM"

        draw.text((36, 32), status_text, fill=status_color)

        # 6. Save annotated photo to class-wise directories
        save_name = f"inspected_{img_path.name}"
        saved_folders = []

        # Always save to master 'all' archive
        draw_img.save(out_root / "all" / save_name, quality=95)

        if num_found == 0:
            draw_img.save(out_root / "clean_no_defects" / save_name, quality=95)
            saved_folders.append("clean_no_defects")
        else:
            if len(detected_classes) > 1:
                draw_img.save(out_root / "multi_defect" / save_name, quality=95)
                saved_folders.append("multi_defect")

            for c_name in detected_classes:
                dest_dir = out_root / c_name
                draw_img.save(dest_dir / save_name, quality=95)
                saved_destinations[c_name] = saved_destinations.get(c_name, 0) + 1
                saved_folders.append(c_name)

        # 7. Build Inspection JSON Report
        img_report = {
            "status": "success",
            "inspection_id": f"insp_{uuid.uuid4().hex[:12]}",
            "timestamp_utc": datetime.now(timezone.utc).isoformat(),
            "model_metadata": {
                "architecture": "Custom Mathematical ML (RawCustomDetector)",
                "checkpoint": str(Path(weights).name),
                "confidence_threshold": conf,
                "nms_iou_threshold": iou,
                "input_resolution": [imgsz, imgsz]
            },
            "image_metadata": {
                "filename": img_path.name,
                "original_path": str(img_path).replace("\\", "/"),
                "file_size_bytes": file_size,
                "dimensions": {
                    "width": orig_w,
                    "height": orig_h,
                    "channels": 3
                }
            },
            "summary": {
                "total_defects": num_found,
                "overall_severity": overall_severity,
                "is_clean": num_found == 0,
                "is_multi_defect": len(detected_classes) > 1,
                "detected_classes": sorted(list(detected_classes)),
                "primary_defect": list(detected_classes)[0] if detected_classes else None,
                "severity_breakdown": {k: v for k, v in severity_counts.items() if v > 0}
            },
            "detections": detections_payload,
            "storage_artifacts": {
                "annotated_filename": save_name,
                "primary_folder": saved_folders[0] if saved_folders else "clean_no_defects",
                "all_saved_folders": saved_folders
            }
        }

        # Save single image JSON
        single_json_path = out_root / "json" / f"{img_path.stem}.json"
        with open(single_json_path, "w", encoding="utf-8") as jf:
            json.dump(img_report, jf, indent=2)

        batch_json_reports.append(img_report)

    total_latency_ms = (time.time() - start_total_time) * 1000.0

    # 8. Save consolidated batch inspection report
    consolidated_report = {
        "status": "success",
        "batch_id": f"batch_{uuid.uuid4().hex[:12]}",
        "generated_at_utc": datetime.now(timezone.utc).isoformat(),
        "summary": {
            "total_images_inspected": len(image_paths),
            "total_defects_found": total_defects,
            "defect_class_distribution": saved_destinations,
            "average_latency_ms": round(total_latency_ms / max(1, len(image_paths)), 2)
        },
        "inspections": batch_json_reports
    }

    consolidated_json_path = out_root / "inspection_report.json"
    with open(consolidated_json_path, "w", encoding="utf-8") as cjf:
        json.dump(consolidated_report, cjf, indent=2)

    print("\n" + "=" * 75)
    print(f"🎉 Detection complete! Found {total_defects} defect(s) across {len(image_paths)} image(s).")
    print(f"📁 Class-Wise Organization at: {output_dir}")
    for c_name, count in saved_destinations.items():
        if count > 0:
            print(f"    - {c_name:<18}: {count} image(s)")
    print(f"\n💾 Backend Database JSON Outputs:")
    print(f"    - Single JSONs:       {out_root / 'json'}/<image_name>.json")
    print(f"    - Consolidated Batch: {consolidated_json_path}")
    print("=" * 75)

    return consolidated_report


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Run Custom Mathematical ML Object Detection Inference")
    parser.add_argument("--source", type=str, required=True, help="Path to image file or directory")
    parser.add_argument("--weights", type=str, default="d:/DETECTION_MODEL/runs/detect/custom_tower_defects/weights/best.pt", help="Path to weights .pt")
    parser.add_argument("--conf", type=float, default=0.25, help="Confidence threshold")
    parser.add_argument("--iou", type=float, default=0.45, help="NMS IoU threshold")
    parser.add_argument("--imgsz", type=int, default=640, help="Image resolution")
    parser.add_argument("--device", type=str, default="auto", help="Compute device ('cpu', 'cuda', 'auto')")
    parser.add_argument("--output", type=str, default="d:/DETECTION_MODEL/runs/detect/predictions", help="Output directory")

    args = parser.parse_args()
    run_custom_prediction(
        source=args.source,
        weights=args.weights,
        conf=args.conf,
        iou=args.iou,
        imgsz=args.imgsz,
        device=args.device,
        output_dir=args.output
    )