src/utils/visualize_all.py

# utils/visualize_all.py

import cv2
import numpy as np

def check_collision(box1, box2):
    x1_1, y1_1, x2_1, y2_1 = box1
    x1_2, y1_2, x2_2, y2_2 = box2
    return not (x2_1 < x1_2 or x2_2 < x1_1 or y2_1 < y1_2 or y2_2 < y1_1)

def find_non_overlapping_position(bg_box, existing_boxes, image_shape, bbox):
    x1, y1, x2, y2 = bbox
    bg_w = bg_box[2] - bg_box[0]
    bg_h = bg_box[3] - bg_box[1]

    positions = [
        (x2 + 10, y1, "right"),
        (x1 - bg_w - 10, y1, "left"),
        (x1, y1 - bg_h - 10, "top"),
        (x1, y2 + 10, "bottom"),
        (x2 + 10, y1 - bg_h - 10, "top-right"),
        (x2 + 10, y2 + 10, "bottom-right"),
        (x1 - bg_w - 10, y1 - bg_h - 10, "top-left"),
        (x1 - bg_w - 10, y2 + 10, "bottom-left"),
    ]

    for pos_x, pos_y, side in positions:
        if pos_x < 10:
            pos_x = 10
        if pos_x + bg_w > image_shape[1] - 10:
            pos_x = image_shape[1] - bg_w - 10
        if pos_y < 10:
            pos_y = 10
        if pos_y + bg_h > image_shape[0] - 10:
            pos_y = image_shape[0] - bg_h - 10

        new_box = (pos_x, pos_y, pos_x + bg_w, pos_y + bg_h)
        collision = False
        for existing_box in existing_boxes:
            if check_collision(new_box, existing_box):
                collision = True
                break
        if not collision:
            return (pos_x, pos_y), side
    return (x2 + 10, y1), "right"

def visualize_detections(image, detections, masks, measurements):
    overlay = image.copy()
    info_boxes = []

    # First pass: build info boxes
    for i, det in enumerate(detections):
        label = det['label']
        x1, y1, x2, y2 = map(int, det['bbox'])
        dims = measurements[i] if i < len(measurements) else {}

        if dims:
            dim_texts = []
            cls = label.lower()

            if cls == "bolt":
                nominal = dims.get("Nominal_M", dims.get("Nominal_Dia", dims.get("Nominal Dia", "")))
                length = dims.get("Length_mm", dims.get("Length", None))
                if nominal:
                    # Convert string or numeric to M-format
                    if isinstance(nominal, str):
                        try:
                            nominal_float = float(nominal)
                            if nominal_float > 0:
                                nominal = f"M{int(nominal_float)}"
                        except (ValueError, TypeError):
                            pass  # Keep as-is if not numeric
                    elif isinstance(nominal, (int, float)) and nominal > 0:
                        nominal = f"M{int(nominal)}"
                    dim_texts.append(f"bolt size = {nominal}")
                if length is not None:
                    # Remove decimal .0 and no space before mm
                    len_str = f"{length:.0f}" if length == int(length) else f"{length:.1f}"
                    dim_texts.append(f"Length = {len_str}mm")

            elif cls == "nut":
                nominal = dims.get("Nominal_Dia", dims.get("Nominal_M", dims.get("Nominal Dia", "")))
                af = dims.get("AF", dims.get("AF_mm", None))
                if nominal:
                    # Convert string or numeric to M-format
                    if isinstance(nominal, str):
                        try:
                            nominal_float = float(nominal)
                            if nominal_float > 0:
                                nominal = f"M{int(nominal_float)}"
                        except (ValueError, TypeError):
                            pass
                    elif isinstance(nominal, (int, float)) and nominal > 0:
                        nominal = f"M{int(nominal)}"
                    dim_texts.append(f"nut size = {nominal}")
                if af is not None:
                    af_str = f"{af:.0f}" if af == int(af) else f"{af:.1f}"
                    dim_texts.append(f"AF = {af_str}mm")

            elif cls == "washer":
                nominal = dims.get("Nominal_M", dims.get("Nominal_Dia", dims.get("Nominal Dia", "")))
                od = dims.get("OD", dims.get("OD_mm", None))
                id_val = dims.get("ID", dims.get("ID_mm", None))
                # Convert string or numeric to M-format
                if nominal:
                    if isinstance(nominal, str):
                        try:
                            nominal_float = float(nominal)
                            if nominal_float > 0:
                                nominal = f"M{int(nominal_float)}"
                        except (ValueError, TypeError):
                            pass
                    elif isinstance(nominal, (int, float)) and nominal > 0:
                        nominal = f"M{int(nominal)}"
                # First line: M6,OD=10mm
                first_line = nominal
                if od is not None:
                    od_str = f"{od:.0f}" if od == int(od) else f"{od:.1f}"
                    first_line += f",OD={od_str}mm"
                dim_texts.append(first_line)
                # Second line: ID=5mm
                if id_val is not None:
                    id_str = f"{id_val:.0f}" if id_val == int(id_val) else f"{id_val:.1f}"
                    dim_texts.append(f"ID={id_str}mm")

            elif cls == "screw":
                nominal = dims.get("Nominal_Dia", dims.get("Nominal_M", dims.get("Nominal Dia", "")))
                length = dims.get("Length_mm", dims.get("Length", None))
                if nominal:
                    # Convert string or numeric to M-format
                    if isinstance(nominal, str):
                        try:
                            nominal_float = float(nominal)
                            if nominal_float > 0:
                                nominal = f"M{int(nominal_float)}"
                        except (ValueError, TypeError):
                            pass
                    elif isinstance(nominal, (int, float)) and nominal > 0:
                        nominal = f"M{int(nominal)}"
                    dim_texts.append(f"screw size = {nominal}")
                if length is not None:
                    len_str = f"{length:.0f}" if length == int(length) else f"{length:.1f}"
                    dim_texts.append(f"Length = {len_str}mm")

            # Build text box dimensions
            if dim_texts:
                max_text_width = 0
                total_text_height = len(dim_texts) * 25
                for txt in dim_texts:
                    (tw, th), _ = cv2.getTextSize(txt, cv2.FONT_HERSHEY_SIMPLEX, 0.7, 2)
                    max_text_width = max(max_text_width, tw)

                bg_w = max_text_width + 10
                bg_h = total_text_height + 10
                existing_boxes = [box[:4] for box in info_boxes]
                (text_x, text_y), side = find_non_overlapping_position(
                    (0, 0, bg_w, bg_h), existing_boxes, image.shape, (x1, y1, x2, y2)
                )
                bg_x1 = text_x - 5
                bg_y1 = text_y - 5
                bg_x2 = text_x + max_text_width + 5
                bg_y2 = text_y + total_text_height + 5
                info_boxes.append((bg_x1, bg_y1, bg_x2, bg_y2, text_x, text_y, dim_texts, (x1, y1, x2, y2), side))

    # Second pass: draw masks, bbox, labels, info boxes
    for i, det in enumerate(detections):
        label = det['label']
        x1, y1, x2, y2 = map(int, det['bbox'])
        dims = measurements[i] if i < len(measurements) else {}

        # Mask
        if i < len(masks) and masks[i].max() > 0:
            mask = masks[i].astype(np.uint8)
            if mask.max() > 1:
                mask = (mask > 127).astype(np.uint8)
            colored_mask = np.zeros_like(image, dtype=np.uint8)
            colored_mask[:, :, 2] = mask * 255
            overlay = np.where(mask[:, :, None].astype(bool),
                               cv2.addWeighted(overlay, 0.7, colored_mask, 0.3, 0),
                               overlay)

        # Bounding box
        cv2.rectangle(overlay, (x1, y1), (x2, y2), (255, 255, 255), 2)

        # Object label
        label_text = f"{label}"
        (tw, th), baseline = cv2.getTextSize(label_text, cv2.FONT_HERSHEY_SIMPLEX, 0.7, 2)
        cv2.rectangle(overlay, (x1, y1 - th - 8), (x1 + tw + 6, y1), (50, 50, 50), -1)
        cv2.putText(overlay, label_text, (x1 + 3, y1 - 5),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 255), 2)

    # Draw info boxes
    for (bg_x1, bg_y1, bg_x2, bg_y2, text_x, text_y, dim_texts, bbox, side) in info_boxes:
        x1, y1, x2, y2 = bbox
        box_center_x = (bg_x1 + bg_x2) // 2
        box_center_y = (bg_y1 + bg_y2) // 2
        obj_center_x = (x1 + x2) // 2
        obj_center_y = (y1 + y2) // 2
        distance = np.sqrt((box_center_x - obj_center_x)**2 + (box_center_y - obj_center_y)**2)

        if distance > 50:
            if side == "right":
                line_start = (x2, (y1 + y2) // 2)
                line_end = (bg_x1, box_center_y)
            elif side == "left":
                line_start = (x1, (y1 + y2) // 2)
                line_end = (bg_x2, box_center_y)
            elif side == "top":
                line_start = ((x1 + x2) // 2, y1)
                line_end = (box_center_x, bg_y2)
            elif side == "bottom":
                line_start = ((x1 + x2) // 2, y2)
                line_end = (box_center_x, bg_y1)
            else:
                line_start = (x2, y1) if "right" in side else (x1, y1)
                line_end = (box_center_x, box_center_y)
            cv2.line(overlay, line_start, line_end, (0, 255, 0), 2)
            cv2.circle(overlay, line_start, 3, (0, 255, 0), -1)
            cv2.circle(overlay, line_end, 3, (0, 255, 0), -1)

        cv2.rectangle(overlay, (bg_x1, bg_y1), (bg_x2, bg_y2), (0, 0, 0), -1)
        cv2.rectangle(overlay, (bg_x1, bg_y1), (bg_x2, bg_y2), (0, 150, 0), 2)

        for j, txt in enumerate(dim_texts):
            text_y_pos = text_y + 25 + j * 25
            cv2.putText(overlay, txt, (text_x, text_y_pos),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 255), 2)

    return overlay