msintui/Intelligent_PID
0
1import os2import math3import torch4import cv25import numpy as np6from typing import List, Optional, Tuple, Dict7from dataclasses import replace8from math import sqrt9import json10import uuid11from pathlib import Path12 13# Base classes and utilities14from base import BaseDetector15from detection_schema import DetectionContext16from utils import DebugHandler17from config import SymbolConfig, TagConfig, LineConfig, PointConfig, JunctionConfig18 19# DeepLSD model for line detection20from deeplsd.models.deeplsd_inference import DeepLSD21from ultralytics import YOLO22 23# Detection schema: dataclasses for different objects24from detection_schema import (25 BBox,26 Coordinates,27 Point,28 Line,29 Symbol,30 Tag,31 SymbolType,32 LineStyle,33 ConnectionType,34 JunctionType,35 Junction36)37 38# Skeletonization and label processing for junction detection39from skimage.morphology import skeletonize40from skimage.measure import label41 42 43import os44import cv245import torch46import numpy as np47from dataclasses import replace48from typing import List, Optional49from detection_utils import robust_merge_lines50 51 52class LineDetector(BaseDetector):53 """54 DeepLSD-based line detection with patch-based tiling and global merging.55 """56 57 def __init__(self,58 config: LineConfig,59 model_path: str,60 model_config: dict,61 device: torch.device,62 debug_handler: DebugHandler = None):63 super().__init__(config, debug_handler)64 65 # Fix device selection for Apple Silicon66 if torch.backends.mps.is_available():67 self.device = torch.device("mps")68 elif torch.cuda.is_available():69 self.device = torch.device("cuda")70 else:71 self.device = torch.device("cpu")72 73 self.model_path = model_path74 self.model_config = model_config75 self.model = self._load_model(model_path)76 77 # Patch parameters78 self.patch_size = 51279 self.overlap = 1080 81 # Merging thresholds82 self.angle_thresh = 5.0 # degrees83 self.dist_thresh = 5.0 # pixels84 85 def _preprocess(self, image: np.ndarray) -> np.ndarray:86 kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (2, 2))87 dilated = cv2.dilate(image, kernel, iterations=2)88 89 skeleton = cv2.bitwise_not(dilated)90 skeleton = skeletonize(skeleton // 255)91 skeleton = (skeleton * 255).astype(np.uint8)92 kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (1, 1))93 clean_image = cv2.dilate(skeleton, kernel, iterations=5)94 95 self.debug_handler.save_artifact(name="skeleton", data=clean_image, extension="png")96 97 return clean_image98 99 def _postprocess(self, image: np.ndarray) -> np.ndarray:100 return None101 # -------------------------------------102 # 1) Load Model103 # -------------------------------------104 def _load_model(self, model_path: str) -> DeepLSD:105 if not os.path.exists(model_path):106 raise FileNotFoundError(f"Model file not found: {model_path}")107 ckpt = torch.load(model_path, map_location=self.device)108 model = DeepLSD(self.model_config)109 model.load_state_dict(ckpt["model"])110 return model.to(self.device).eval()111 112 # -------------------------------------113 # 2) Main Detection Pipeline114 # -------------------------------------115 def detect(self,116 image: np.ndarray,117 context: DetectionContext,118 mask_coords: Optional[List[BBox]] = None,119 *args,120 **kwargs) -> None:121 """122 Steps:123 - Optional mask + threshold124 - Tile into overlapping patches125 - For each patch => run DeepLSD => re-map lines to global coords126 - Merge lines robustly127 - Build final Line objects => add to context128 """129 mask_coords = mask_coords or []130 131 skeleton = self._preprocess(image)132 # (A) Optional mask + threshold if you want a binary133 # If your model expects grayscale or binary, do it here:134 processed_img = self._apply_mask_and_threshold(skeleton, mask_coords)135 # (B) Patch-based inference => collect raw lines in global coords136 all_lines = self._detect_in_patches(processed_img)137 138 # (C) Merge the lines in the global coordinate system139 merged_line_segments = robust_merge_lines(140 all_lines,141 angle_thresh=self.angle_thresh,142 dist_thresh=self.dist_thresh143 )144 145 # (D) Convert merged segments => final Line objects, add to context146 for (x1, y1, x2, y2) in merged_line_segments:147 line_obj = self._create_line_object(x1, y1, x2, y2)148 context.add_line(line_obj)149 150 # -------------------------------------151 # 3) Optional Mask + Threshold152 # -------------------------------------153 def _apply_mask_and_threshold(self, image: np.ndarray, mask_coords: List[BBox]) -> np.ndarray:154 """White out rectangular areas, then threshold to binary (if needed)."""155 masked = image.copy()156 for bbox in mask_coords:157 x1, y1 = int(bbox.xmin), int(bbox.ymin)158 x2, y2 = int(bbox.xmax), int(bbox.ymax)159 cv2.rectangle(masked, (x1, y1), (x2, y2), (255, 255, 255), -1)160 161 # If image has 3 channels, convert to grayscale162 if len(masked.shape) == 3:163 masked_gray = cv2.cvtColor(masked, cv2.COLOR_BGR2GRAY)164 else:165 masked_gray = masked166 167 # Binary threshold (adjust threshold as needed)168 # If your model expects a plain grayscale, skip threshold169 binary_img = cv2.threshold(masked_gray, 127, 255, cv2.THRESH_BINARY)[1]170 return binary_img171 172 # -------------------------------------173 # 4) Patch-Based Inference174 # -------------------------------------175 def _detect_in_patches(self, processed_img: np.ndarray) -> List[tuple]:176 """177 Break the image into overlapping patches, run DeepLSD,178 map local lines => global coords, and return the global line list.179 """180 patch_size = self.patch_size181 overlap = self.overlap182 183 height, width = processed_img.shape[:2]184 step = patch_size - overlap185 186 all_lines = []187 188 for y in range(0, height, step):189 patch_ymax = min(y + patch_size, height)190 patch_ymin = patch_ymax - patch_size if (patch_ymax - y) < patch_size else y191 if patch_ymin < 0: patch_ymin = 0192 193 for x in range(0, width, step):194 patch_xmax = min(x + patch_size, width)195 patch_xmin = patch_xmax - patch_size if (patch_xmax - x) < patch_size else x196 if patch_xmin < 0: patch_xmin = 0197 198 patch = processed_img[patch_ymin:patch_ymax, patch_xmin:patch_xmax]199 200 # Run model201 local_lines = self._run_model_inference(patch)202 203 # Convert local lines => global coords204 for ln in local_lines:205 (x1_local, y1_local), (x2_local, y2_local) = ln206 207 # offset by patch_xmin, patch_ymin208 gx1 = x1_local + patch_xmin209 gy1 = y1_local + patch_ymin210 gx2 = x2_local + patch_xmin211 gy2 = y2_local + patch_ymin212 213 # Optional: clamp or filter lines partially out-of-bounds214 if 0 <= gx1 < width and 0 <= gx2 < width and 0 <= gy1 < height and 0 <= gy2 < height:215 all_lines.append((gx1, gy1, gx2, gy2))216 217 return all_lines218 219 # -------------------------------------220 # 5) Model Inference (Single Patch)221 # -------------------------------------222 def _run_model_inference(self, patch_img: np.ndarray) -> np.ndarray:223 """224 Run DeepLSD on a single patch (already masked/thresholded).225 patch_img shape: [patchH, patchW].226 Returns lines shape: [N, 2, 2].227 """228 # Convert patch to float32 and scale229 inp = torch.tensor(patch_img, dtype=torch.float32, device=self.device)[None, None] / 255.0230 with torch.no_grad():231 output = self.model({"image": inp})232 lines = output["lines"][0] # shape (N, 2, 2)233 return lines234 235 # -------------------------------------236 # 6) Convert Merged Segments => Line Objects237 # -------------------------------------238 def _create_line_object(self, x1: float, y1: float, x2: float, y2: float) -> Line:239 """240 Create a minimal `Line` object from the final merged coordinates.241 """242 margin = 2243 # Start point244 start_pt = Point(245 coords=Coordinates(int(x1), int(y1)),246 bbox=BBox(247 xmin=int(x1 - margin),248 ymin=int(y1 - margin),249 xmax=int(x1 + margin),250 ymax=int(y1 + margin)251 ),252 type=JunctionType.END,253 confidence=1.0254 )255 # End point256 end_pt = Point(257 coords=Coordinates(int(x2), int(y2)),258 bbox=BBox(259 xmin=int(x2 - margin),260 ymin=int(y2 - margin),261 xmax=int(x2 + margin),262 ymax=int(y2 + margin)263 ),264 type=JunctionType.END,265 confidence=1.0266 )267 268 # Overall bounding box269 x_min = int(min(x1, x2))270 x_max = int(max(x1, x2))271 y_min = int(min(y1, y2))272 y_max = int(max(y1, y2))273 274 line_obj = Line(275 start=start_pt,276 end=end_pt,277 bbox=BBox(xmin=x_min, ymin=y_min, xmax=x_max, ymax=y_max),278 style=LineStyle(279 connection_type=ConnectionType.SOLID,280 stroke_width=2,281 color="#000000"282 ),283 confidence=0.9,284 topological_links=[]285 )286 return line_obj287 288class PointDetector(BaseDetector):289 """290 A detector that:291 1) Reads lines from the context292 2) Clusters endpoints within 'threshold_distance'293 3) Updates lines so that shared endpoints reference the same Point object294 """295 296 def __init__(self,297 config:PointConfig,298 debug_handler: DebugHandler = None):299 super().__init__(config, debug_handler) # No real model to load300 self.threshold_distance = config.threshold_distance301 302 def _load_model(self, model_path: str):303 """No model needed for simple point unification."""304 return None305 306 def detect(self, image: np.ndarray, context: DetectionContext, *args, **kwargs) -> None:307 """308 Main method called by the pipeline.309 1) Gather all line endpoints from context310 2) Cluster them within 'threshold_distance'311 3) Update the line endpoints so they reference the unified cluster point312 """313 # 1) Collect all endpoints314 endpoints = []315 for line in context.lines.values():316 endpoints.append(line.start)317 endpoints.append(line.end)318 319 # 2) Cluster endpoints320 clusters = self._cluster_points(endpoints, self.threshold_distance)321 322 # 3) Build a dictionary of "representative" points323 # So that each cluster has one "canonical" point324 # Then we link all the points in that cluster to the canonical reference325 unified_point_map = {}326 for cluster in clusters:327 # let's pick the first point in the cluster as the "representative"328 rep_point = cluster[0]329 for p in cluster[1:]:330 unified_point_map[p.id] = rep_point331 332 # 4) Update all lines to reference the canonical point333 for line in context.lines.values():334 # unify start335 if line.start.id in unified_point_map:336 line.start = unified_point_map[line.start.id]337 # unify end338 if line.end.id in unified_point_map:339 line.end = unified_point_map[line.end.id]340 341 # We could also store the final set of unique points back in context.points342 # (e.g. clearing old duplicates).343 # That step is optional: you might prefer to keep everything in lines only,344 # or you might want context.points as a separate reference.345 346 # If you want to keep unique points in context.points:347 new_points = {}348 for line in context.lines.values():349 new_points[line.start.id] = line.start350 new_points[line.end.id] = line.end351 context.points = new_points # replace the dictionary of points352 353 def _preprocess(self, image: np.ndarray) -> np.ndarray:354 """No specific image preprocessing needed."""355 return image356 357 def _postprocess(self, image: np.ndarray) -> np.ndarray:358 """No specific image postprocessing needed."""359 return image360 361 # ----------------------362 # HELPER: clustering363 # ----------------------364 def _cluster_points(self, points: List[Point], threshold: float) -> List[List[Point]]:365 """366 Very naive clustering:367 1) Start from the first point368 2) If it's within threshold of an existing cluster's representative,369 put it in that cluster370 3) Otherwise start a new cluster371 Return: list of clusters, each is a list of Points372 """373 clusters = []374 375 for pt in points:376 placed = False377 for cluster in clusters:378 # pick the first point in the cluster as reference379 ref_pt = cluster[0]380 if self._distance(pt, ref_pt) < threshold:381 cluster.append(pt)382 placed = True383 break384 385 if not placed:386 clusters.append([pt])387 388 return clusters389 390 def _distance(self, p1: Point, p2: Point) -> float:391 dx = p1.coords.x - p2.coords.x392 dy = p1.coords.y - p2.coords.y393 return sqrt(dx*dx + dy*dy)394 395 396class JunctionDetector(BaseDetector):397 """398 Classifies points as 'END', 'L', or 'T' by skeletonizing the binarized image399 and analyzing local connectivity. Also creates Junction objects in the context.400 """401 402 def __init__(self, config: JunctionConfig, debug_handler: DebugHandler = None):403 super().__init__(config, debug_handler) # no real model path404 self.window_size = config.window_size405 self.radius = config.radius406 self.angle_threshold_lb = config.angle_threshold_lb407 self.angle_threshold_ub = config.angle_threshold_ub408 self.debug_handler = debug_handler or DebugHandler()409 410 def _load_model(self, model_path: str):411 """Not loading any actual model, just skeleton logic."""412 return None413 414 def detect(self,415 image: np.ndarray,416 context: DetectionContext,417 *args,418 **kwargs) -> None:419 """420 1) Convert to binary & skeletonize421 2) Classify each point in the context422 3) Create a Junction for each point and store it in context.junctions423 (with 'connected_lines' referencing lines that share this point).424 """425 # 1) Preprocess -> skeleton426 skeleton = self._create_skeleton(image)427 428 # 2) Classify each point429 for pt in context.points.values():430 pt.type = self._classify_point(skeleton, pt)431 432 # 3) Create a Junction object for each point433 # If you prefer only T or L, you can filter out END points.434 self._record_junctions_in_context(context)435 436 def _preprocess(self, image: np.ndarray) -> np.ndarray:437 """We might do thresholding; let's do a simple binary threshold."""438 if image.ndim == 3:439 gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)440 else:441 gray = image442 _, bin_image = cv2.threshold(gray, 127, 255, cv2.THRESH_BINARY)443 return bin_image444 445 def _postprocess(self, image: np.ndarray) -> np.ndarray:446 return image447 448 def _create_skeleton(self, raw_image: np.ndarray) -> np.ndarray:449 """Skeletonize the binarized image."""450 bin_img = self._preprocess(raw_image)451 # For skeletonize, we need a boolean array452 inv = cv2.bitwise_not(bin_img)453 inv_bool = (inv > 127).astype(np.uint8)454 skel = skeletonize(inv_bool).astype(np.uint8) * 255455 return skel456 457 def _classify_point(self, skeleton: np.ndarray, pt: Point) -> JunctionType:458 """459 Given a skeleton image, look around 'pt' in a local window460 to determine if it's an END, L, or T.461 """462 classification = JunctionType.END # default463 464 half_w = self.window_size // 2465 x, y = pt.coords.x, pt.coords.y466 467 top = max(0, y - half_w)468 bottom = min(skeleton.shape[0], y + half_w + 1)469 left = max(0, x - half_w)470 right = min(skeleton.shape[1], x + half_w + 1)471 472 patch = (skeleton[top:bottom, left:right] > 127).astype(np.uint8)473 474 # create circular mask475 circle_mask = np.zeros_like(patch, dtype=np.uint8)476 local_cx = x - left477 local_cy = y - top478 cv2.circle(circle_mask, (local_cx, local_cy), self.radius, 1, -1)479 circle_skel = patch & circle_mask480 481 # label connected regions482 labeled = label(circle_skel, connectivity=2)483 num_exits = labeled.max()484 485 if num_exits == 1:486 classification = JunctionType.END487 elif num_exits == 2:488 # check angle for L489 classification = self._check_angle_for_L(labeled)490 elif num_exits == 3:491 classification = JunctionType.T492 493 return classification494 495 def _check_angle_for_L(self, labeled_region: np.ndarray) -> JunctionType:496 """497 If the angle between two branches is within498 [angle_threshold_lb, angle_threshold_ub], it's 'L'.499 Otherwise default to END.500 """501 coords = np.argwhere(labeled_region == 1)502 if len(coords) < 2:503 return JunctionType.END504 505 (y1, x1), (y2, x2) = coords[:2]506 dx = x2 - x1507 dy = y2 - y1508 angle = math.degrees(math.atan2(dy, dx))509 acute_angle = min(abs(angle), 180 - abs(angle))510 511 if self.angle_threshold_lb <= acute_angle <= self.angle_threshold_ub:512 return JunctionType.L513 return JunctionType.END514 515 # -----------------------------------------516 # EXTRA STEP: Create Junction objects517 # -----------------------------------------518 def _record_junctions_in_context(self, context: DetectionContext):519 """520 Create a Junction object for each point in context.points.521 If you only want T/L points as junctions, filter them out.522 Also track any lines that connect to this point.523 """524 525 for pt in context.points.values():526 # If you prefer to store all points as junction, do it:527 # or if you want only T or L, do:528 # if pt.type in {JunctionType.T, JunctionType.L}: ...529 530 jn = Junction(531 center=pt.coords,532 junction_type=pt.type,533 # add more properties if needed534 )535 536 # find lines that connect to this point537 connected_lines = []538 for ln in context.lines.values():539 if ln.start.id == pt.id or ln.end.id == pt.id:540 connected_lines.append(ln.id)541 542 jn.connected_lines = connected_lines543 544 # add to context545 context.add_junction(jn)546 547import json548import uuid549 550class SymbolDetector(BaseDetector):551 """552 A placeholder detector that reads precomputed symbol data553 from a JSON file and populates the context with Symbol objects.554 """555 556 def __init__(self,557 config: SymbolConfig,558 debug_handler: Optional[DebugHandler] = None,559 symbol_json_path: str = "./symbols.json"):560 super().__init__(config=config, debug_handler=debug_handler)561 self.symbol_json_path = symbol_json_path562 563 def _load_model(self, model_path: str):564 """Not loading an actual model; symbol data is read from JSON."""565 return None566 567 def detect(self,568 image: np.ndarray,569 context: DetectionContext,570 # roi_offset: Tuple[int, int],571 *args,572 **kwargs) -> None:573 """574 Reads from a JSON file containing symbol info,575 adjusts coordinates using roi_offset, and updates context.576 """577 symbol_data = self._load_json_data(self.symbol_json_path)578 if not symbol_data:579 return580 581 # x_min, y_min = roi_offset # Offset values from cropping582 583 for record in symbol_data.get("detections", []): # Fix: Use "detections" key584 # sym_obj = self._parse_symbol_record(record, x_min, y_min)585 sym_obj = self._parse_symbol_record(record)586 context.add_symbol(sym_obj)587 588 def _preprocess(self, image: np.ndarray) -> np.ndarray:589 return image590 591 def _postprocess(self, image: np.ndarray) -> np.ndarray:592 return image593 594 # --------------595 # HELPER METHODS596 # --------------597 def _load_json_data(self, json_path: str) -> dict:598 if not os.path.exists(json_path):599 self.debug_handler.save_artifact(name="symbol_error",600 data=b"Missing symbol JSON file",601 extension="txt")602 return {}603 604 with open(json_path, "r", encoding="utf-8") as f:605 return json.load(f)606 607 def _parse_symbol_record(self, record: dict) -> Symbol:608 """609 Builds a Symbol object from a JSON record, adjusting coordinates for cropping.610 """611 bbox_list = record.get("bbox", [0, 0, 0, 0])612 # bbox_obj = BBox(613 # xmin=bbox_list[0] - x_min,614 # ymin=bbox_list[1] - y_min,615 # xmax=bbox_list[2] - x_min,616 # ymax=bbox_list[3] - y_min617 # )618 619 bbox_obj = BBox(620 xmin=bbox_list[0],621 ymin=bbox_list[1],622 xmax=bbox_list[2],623 ymax=bbox_list[3]624 )625 626 627 # Compute the center628 center_coords = Coordinates(629 x=(bbox_obj.xmin + bbox_obj.xmax) // 2,630 y=(bbox_obj.ymin + bbox_obj.ymax) // 2631 )632 633 return Symbol(634 id=record.get("symbol_id", ""),635 class_id=record.get("class_id", -1),636 original_label=record.get("original_label", ""),637 category=record.get("category", ""),638 type=record.get("type", ""),639 label=record.get("label", ""),640 bbox=bbox_obj,641 center=center_coords,642 confidence=record.get("confidence", 0.95),643 model_source=record.get("model_source", ""),644 connections=[]645 )646 647class TagDetector(BaseDetector):648 """649 A placeholder detector that reads precomputed tag data650 from a JSON file and populates the context with Tag objects.651 """652 653 def __init__(self,654 config: TagConfig,655 debug_handler: Optional[DebugHandler] = None,656 tag_json_path: str = "./tags.json"):657 super().__init__(config=config, debug_handler=debug_handler)658 self.tag_json_path = tag_json_path659 660 def _load_model(self, model_path: str):661 """Not loading an actual model; tag data is read from JSON."""662 return None663 664 def detect(self,665 image: np.ndarray,666 context: DetectionContext,667 # roi_offset: Tuple[int, int],668 *args,669 **kwargs) -> None:670 """671 Reads from a JSON file containing tag info,672 adjusts coordinates using roi_offset, and updates context.673 """674 675 tag_data = self._load_json_data(self.tag_json_path)676 if not tag_data:677 return678 679 # x_min, y_min = roi_offset # Offset values from cropping680 681 for record in tag_data.get("detections", []): # Fix: Use "detections" key682 # tag_obj = self._parse_tag_record(record, x_min, y_min)683 tag_obj = self._parse_tag_record(record)684 context.add_tag(tag_obj)685 686 def _preprocess(self, image: np.ndarray) -> np.ndarray:687 return image688 689 def _postprocess(self, image: np.ndarray) -> np.ndarray:690 return image691 692 # --------------693 # HELPER METHODS694 # --------------695 def _load_json_data(self, json_path: str) -> dict:696 if not os.path.exists(json_path):697 self.debug_handler.save_artifact(name="tag_error",698 data=b"Missing tag JSON file",699 extension="txt")700 return {}701 702 with open(json_path, "r", encoding="utf-8") as f:703 return json.load(f)704 705 def _parse_tag_record(self, record: dict) -> Tag:706 """707 Builds a Tag object from a JSON record, adjusting coordinates for cropping.708 """709 bbox_list = record.get("bbox", [0, 0, 0, 0])710 # bbox_obj = BBox(711 # xmin=bbox_list[0] - x_min,712 # ymin=bbox_list[1] - y_min,713 # xmax=bbox_list[2] - x_min,714 # ymax=bbox_list[3] - y_min715 # )716 717 bbox_obj = BBox(718 xmin=bbox_list[0],719 ymin=bbox_list[1],720 xmax=bbox_list[2],721 ymax=bbox_list[3]722 )723 724 return Tag(725 text=record.get("text", ""),726 bbox=bbox_obj,727 confidence=record.get("confidence", 1.0),728 source=record.get("source", ""),729 text_type=record.get("text_type", "Unknown"),730 id=record.get("id", str(uuid.uuid4())),731 font_size=record.get("font_size", 12),732 rotation=record.get("rotation", 0.0)733 )