app.py

import base64
import shutil
import traceback
from pathlib import Path

import cv2
import gradio as gr
import numpy as np
import yaml
from PIL import Image
from ultralytics import YOLO

from src.utils.detect_objects import run_modelA
from src.utils.generate_masks import run_modelB, seg_model as default_seg_model
from src.utils.reference import detect_reference
from src.measure.measure_tool import process_measurements
from src.utils.match_spec import run_spec_match
from src.utils.visualize_all import visualize_detections

ROOT = Path(__file__).resolve().parent
CONFIG_PATH = ROOT / "config" / "settings.yaml"
MODELS_DIR = ROOT / "models"
DATASETS_DIR = ROOT / "data" / "datasets"
INPUTS_DIR = ROOT / "inputs"
OUTPUTS_DIR = ROOT / "outputs"
DEFAULT_PX_PER_MM = 10.0


def load_settings():
    if not CONFIG_PATH.exists():
        raise FileNotFoundError(f"Missing config file: {CONFIG_PATH}")
    with open(CONFIG_PATH, "r", encoding="utf-8") as f:
        return yaml.safe_load(f)


settings = load_settings()
DETECTION_CFG = settings.get("detection", {})
REFERENCE_CFG = settings.get("reference", {})

REFERENCE_SIZE_MM = float(REFERENCE_CFG.get("size_mm", 20.0))
DETECTION_DEVICE = DETECTION_CFG.get("device", "cpu")
DETECTION_IMAGE_SIZE = int(DETECTION_CFG.get("image_size", 640))
DETECTION_CONF = float(DETECTION_CFG.get("confidence_threshold", 0.25))
DETECTION_IOU = float(DETECTION_CFG.get("iou_threshold", 0.5))

MODEL_A_PATH = MODELS_DIR / "model_a.pt"
MODEL_B_PATH = MODELS_DIR / "model_b.pt"


def ensure_directories():
    for d in [
        INPUTS_DIR,
        OUTPUTS_DIR / "1_captured_images",
        OUTPUTS_DIR / "2_reference",
        OUTPUTS_DIR / "3_detection" / "labels",
        OUTPUTS_DIR / "4_segmentation" / "masks",
        OUTPUTS_DIR / "4_segmentation" / "overlay",
        OUTPUTS_DIR / "5_measured",
        OUTPUTS_DIR / "6_results" / "spec_match_report",
        OUTPUTS_DIR / "6_results" / "output_images",
    ]:
        d.mkdir(parents=True, exist_ok=True)


def clean_measured_folder():
    measured_dir = OUTPUTS_DIR / "5_measured"
    if measured_dir.exists():
        for f in measured_dir.glob("*"):
            if f.is_file():
                f.unlink()


def pil_to_bgr(image):
    if isinstance(image, np.ndarray):
        rgb = image
    elif isinstance(image, Image.Image):
        rgb = np.array(image)
    else:
        raise ValueError("Unsupported image type")
    if rgb.ndim == 2:
        rgb = cv2.cvtColor(rgb, cv2.COLOR_GRAY2RGB)
    if rgb.shape[2] == 4:
        rgb = cv2.cvtColor(rgb, cv2.COLOR_RGBA2RGB)
    return cv2.cvtColor(rgb, cv2.COLOR_RGB2BGR)


def save_input_image(image, filename="input.jpg"):
    img_bgr = pil_to_bgr(image)
    path = INPUTS_DIR / filename
    cv2.imwrite(str(path), img_bgr)
    return path


def encode_image_base64(image):
    success, buffer = cv2.imencode(".jpg", image)
    if not success:
        raise RuntimeError("Could not encode image to JPEG")
    return base64.b64encode(buffer.tobytes()).decode("utf-8")


def build_report_text(spec_results, ref_ok, px_per_mm, fallback_used):
    lines = []
    # Reference status
    if ref_ok:
        lines.append(f"✅ Reference detected → {px_per_mm:.3f} px/mm")
    else:
        if fallback_used:
            lines.append(f"⚠️ Reference NOT found. Using default calibration: {px_per_mm:.3f} px/mm")
        else:
            lines.append("❌ Reference square missing and no fallback active.")

    if not spec_results:
        lines.append("No components measured.")
        return "\n".join(lines)

    # Number of components
    lines.append(f"\nno of components detected - {len(spec_results)}")
    lines.append("=" * 50)

    for idx, res in enumerate(spec_results, 1):
        comp = res['class'].upper()
        measured = res.get('measured', {})
        ref = res.get('reference', {})

        # Determine nominal display
        nominal = "N/A"
        if ref:
            # Try common keys
            nominal = ref.get('Nominal Dia', ref.get('Nominal_M', ref.get('Screw Size', 'N/A')))
        elif measured:
            nominal = measured.get('Nominal_Dia', measured.get('Nominal_M', 'N/A'))

        lines.append(f"{idx}) 🔩 {comp}  |  {nominal}")

        # Measured dimensions (only those that appear in the measured dict)
        if measured:
            lines.append("   Measured values:")
            for k, v in measured.items():
                # Skip internal/verbose keys
                if k in ['class', 'file']:
                    continue
                lines.append(f"      {k}: {v}")
        else:
            lines.append("   No measured dimensions.")

        # Reference (ISO) dimensions
        if ref:
            lines.append("   Standard Specifications (ISO):")
            for k, v in ref.items():
                # Skip keys that are not human‑readable or already shown
                if k in ['Min_Standard_Length', 'Max_Standard_Length', 'Preferred_Standards',
                         'Standard_Length_mm', 'Length_Deviation_mm', 'Length_Valid', 'Length_In_Range',
                         'file', 'class']:
                    continue
                lines.append(f"      {k}: {v}")
        else:
            lines.append("   No ISO match found.")

        lines.append("-" * 40)

    lines.append("=" * 50)
    return "\n".join(lines)


def build_visualization(image_path, detections, mask_array, measurement_results, spec_results):
    image = cv2.imread(str(image_path))
    if image is None:
        raise RuntimeError(f"Unable to read image: {image_path}")

    h, w = image.shape[:2]
    det_list, mask_list, meas_list = [], [], []

    spec_map = {}
    for s in spec_results:
        stem = Path(s["file"]).stem
        # Remove "_measured" suffix to match obj_id format
        key = stem.replace("_measured", "") if "_measured" in stem else stem
        spec_map[key] = s.get("reference", {})

    for idx, det in enumerate(detections):
        det_list.append({"label": det["class_name"], "bbox": det["xyxy"]})

        x1, y1, x2, y2 = map(int, det["xyxy"])
        full_mask = np.zeros((h, w), dtype=np.uint8)
        if mask_array is not None and y2 > y1 and x2 > x1:
            cropped = mask_array[y1:y2, x1:x2]
            if cropped.size > 0:
                mask_crop = (cropped > 127).astype(np.uint8) if cropped.max() > 1 else (cropped > 0).astype(np.uint8)
                full_mask[y1:y2, x1:x2] = mask_crop
        mask_list.append(full_mask)

        cls = det["class_name"].lower()
        obj_id = f"{cls}_{idx+1}"
        meas = measurement_results.get(obj_id, {}).copy()
        matched = spec_map.get(obj_id, {})

        combined = {**meas, **matched}

        # Normalize keys for visualize_all.py
        if cls == "nut":
            # Ensure AF is available
            if "AF" not in combined and "AF_mm" in combined:
                combined["AF"] = combined["AF_mm"]
            # Ensure Nominal_Dia and Nominal_M are set from CSV "Nominal Dia" column
            if "Nominal Dia" in combined:
                if "Nominal_Dia" not in combined:
                    combined["Nominal_Dia"] = combined["Nominal Dia"]
                if "Nominal_M" not in combined:
                    combined["Nominal_M"] = combined["Nominal Dia"]
            # Also handle Nut Size key from CSV
            if "Nut Size" in matched and "Nut Size" not in combined:
                combined["Nut Size"] = matched["Nut Size"]
        elif cls == "washer":
            # Ensure OD and ID are available
            if "OD" not in combined and "OD_mm" in combined:
                combined["OD"] = combined["OD_mm"]
            if "ID" not in combined and "ID_mm" in combined:
                combined["ID"] = combined["ID_mm"]
            # Ensure Nominal_Dia and Nominal_M are set from CSV "Nominal Dia" column
            if "Nominal Dia" in combined:
                if "Nominal_Dia" not in combined:
                    combined["Nominal_Dia"] = combined["Nominal Dia"]
                if "Nominal_M" not in combined:
                    combined["Nominal_M"] = combined["Nominal Dia"]
        elif cls == "bolt":
            # Map CSV column "Bolt Size" to Nominal_M and Nominal_Dia
            if "Bolt Size" in matched:
                combined["Nominal_M"] = matched["Bolt Size"]
                combined["Nominal_Dia"] = matched["Bolt Size"]
            # Also handle "Nominal Dia" if present
            if "Nominal Dia" in combined:
                if "Nominal_M" not in combined:
                    combined["Nominal_M"] = combined["Nominal Dia"]
                if "Nominal_Dia" not in combined:
                    combined["Nominal_Dia"] = combined["Nominal Dia"]
            # Ensure Length_mm
            if "Length_mm" not in combined and "Length" in combined:
                combined["Length_mm"] = combined["Length"]
        elif cls == "screw":
            # Ensure Nominal_Dia and Nominal_M are set from CSV "Nominal Dia" column
            if "Nominal Dia" in combined:
                if "Nominal_Dia" not in combined:
                    combined["Nominal_Dia"] = combined["Nominal Dia"]
                if "Nominal_M" not in combined:
                    combined["Nominal_M"] = combined["Nominal Dia"]
            # Ensure Length_mm is available
            if "Length_mm" not in combined and "Length" in combined:
                combined["Length_mm"] = combined["Length"]
            # Also handle Screw Size key from CSV if needed
            if "Screw Size" in matched and "Screw Size" not in combined:
                combined["Screw Size"] = matched["Screw Size"]

        meas_list.append(combined)

    return visualize_detections(image, det_list, mask_list, meas_list)


def predict(image):
    ensure_directories()
    clean_measured_folder()
    fallback_used = False
    if image is None:
        return {"success": False, "output_image": None, "report_text": "No image provided."}
    try:
        input_path = save_input_image(image, "input.jpg")
        ref_status, px_per_mm, _ = detect_reference(
            image_path=str(input_path),
            ref_size_mm=REFERENCE_SIZE_MM,
            save_path=str(OUTPUTS_DIR / "2_reference"),
        )
        if ref_status != "success":
            px_per_mm = DEFAULT_PX_PER_MM
            fallback_used = True

        if not MODEL_A_PATH.exists() or not MODEL_B_PATH.exists():
            raise FileNotFoundError("YOLO models not found in models/")

        detection_model = YOLO(str(MODEL_A_PATH))
        detections, label_path = run_modelA(
            image_path=str(input_path),
            device=DETECTION_DEVICE,
            imgsz=DETECTION_IMAGE_SIZE,
            conf_thr=DETECTION_CONF,
            iou_thr=DETECTION_IOU,
            save_annotated=True,
            outdir=str(OUTPUTS_DIR / "3_detection"),
            save_labels=True,
            model=detection_model,
        )

        if not detections:
            encoded = encode_image_base64(pil_to_bgr(image))
            return {
                "success": True,
                "output_image": encoded,
                "report_text": "No fasteners detected.",
            }

        mask_array = None
        if label_path:
            mask_result = run_modelB(
                img_path=str(input_path),
                label_txt_path=str(label_path),
                seg_model=default_seg_model,
                masks_dir=str(OUTPUTS_DIR / "4_segmentation" / "masks"),
                images_dir=str(OUTPUTS_DIR / "4_segmentation" / "overlay"),
            )
            mask_array = mask_result.get("mask_array")

        if mask_array is None:
            img_cv = cv2.imread(str(input_path))
            h, w = img_cv.shape[:2]
            mask_array = np.zeros((h, w), dtype=np.uint8)

        measurement_results = process_measurements(
            image_path=str(input_path),
            detections=detections,
            label_path=str(label_path) if label_path else None,
            mask_data=mask_array,
            px_per_mm=px_per_mm,
        )

        spec_results = []
        if measurement_results:
            spec_results = run_spec_match(
                measurements_dir=str(OUTPUTS_DIR / "5_measured"),
                reference_csv_dict={
                    "washer": str(DATASETS_DIR / "washers_dataset.csv"),
                    "bolt":   str(DATASETS_DIR / "bolts_dataset.csv"),
                    "nut":    str(DATASETS_DIR / "nuts_dataset.csv"),
                    "screw":  str(DATASETS_DIR / "screws_dataset.csv"),
                },
                output_txt=str(OUTPUTS_DIR / "6_results" / "spec_match_report" / "spec_match_report.txt"),
            )

        output_image = build_visualization(
            image_path=str(input_path),
            detections=detections,
            mask_array=mask_array,
            measurement_results=measurement_results,
            spec_results=spec_results,
        )

        report = build_report_text(spec_results, ref_status == "success", px_per_mm, fallback_used)

        return {
            "success": True,
            "output_image": encode_image_base64(output_image),
            "report_text": report,
        }
    except Exception as e:
        traceback.print_exc()
        return {"success": False, "output_image": None, "report_text": f"Error: {str(e)}"}


ensure_directories()

with gr.Blocks() as demo:
    gr.Markdown("# 🔧 AFMS — Automated Fastener Measurement System")
    gr.Markdown("Upload an image containing fasteners and a **20mm reference square**.")
    with gr.Row():
        image_input = gr.Image(type="pil", label="Input Image")
    output_json = gr.JSON(label="Results")
    submit = gr.Button("Run Detection", variant="primary")
    submit.click(fn=predict, inputs=[image_input], outputs=[output_json])

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)