# 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