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)