msintui/Intelligent_PID
0
1import cv22import json3import uuid4import os5import logging6from ultralytics import YOLO7from tqdm import tqdm8from storage import StorageInterface9import numpy as np10from typing import Tuple, List, Dict, Any11 12# Configure logging13logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')14 15# Constants16MODEL_PATHS = {17 "model1": "models/Intui_SDM_41.pt",18 "model2": "models/Intui_SDM_20.pt" # Add your second model path here19}20MAX_DIMENSION = 128021CONFIDENCE_THRESHOLDS = [0.1, 0.3, 0.5, 0.7, 0.9]22TEXT_COLOR = (0, 0, 255) # Red color for text23BOX_COLOR = (255, 0, 0) # Red color for box (no transparency)24BG_COLOR = (255, 255, 255, 0.6) # Semi-transparent white for text background25THICKNESS = 1 # Thin text thickness26BOX_THICKNESS = 2 # Box line thickness27MIN_FONT_SCALE = 0.2 # Minimum font scale28MAX_FONT_SCALE = 1.0 # Maximum font scale29TEXT_PADDING = 20 # Increased padding between text elements30OVERLAP_THRESHOLD = 0.3 # Threshold for detecting text overlap31 32def preprocess_image_for_symbol_detection(image_cv: np.ndarray) -> np.ndarray:33 """Preprocess the image for symbol detection."""34 gray = cv2.cvtColor(image_cv, cv2.COLOR_BGR2GRAY)35 equalized = cv2.equalizeHist(gray)36 filtered = cv2.bilateralFilter(equalized, 9, 75, 75)37 edges = cv2.Canny(filtered, 100, 200)38 preprocessed_image = cv2.cvtColor(edges, cv2.COLOR_GRAY2BGR)39 return preprocessed_image40 41def evaluate_detections(detections_list: List[Dict[str, Any]]) -> int:42 """Evaluate the quality of detections."""43 return len(detections_list)44 45def resize_image_with_aspect_ratio(image_cv: np.ndarray, max_dimension: int) -> Tuple[np.ndarray, int, int]:46 """Resize the image while maintaining the aspect ratio."""47 original_height, original_width = image_cv.shape[:2]48 if max(original_width, original_height) > max_dimension:49 scale = max_dimension / float(max(original_width, original_height))50 new_width = int(original_width * scale)51 new_height = int(original_height * scale)52 image_cv = cv2.resize(image_cv, (new_width, new_height), interpolation=cv2.INTER_LINEAR)53 else:54 new_width, new_height = original_width, original_height55 return image_cv, new_width, new_height56 57def merge_detections(all_detections: List[Dict]) -> List[Dict]:58 """59 Merge detections from all models, keeping only the highest confidence detection60 when duplicates are found using IoU.61 """62 if not all_detections:63 return []64 65 # Sort by confidence66 all_detections.sort(key=lambda x: x['confidence'], reverse=True)67 68 # Keep track of which detections to keep69 keep = [True] * len(all_detections)70 71 def calculate_iou(box1, box2):72 """Calculate Intersection over Union (IoU) between two boxes."""73 x1 = max(box1[0], box2[0])74 y1 = max(box1[1], box2[1])75 x2 = min(box1[2], box2[2])76 y2 = min(box1[3], box2[3])77 78 intersection = max(0, x2 - x1) * max(0, y2 - y1)79 area1 = (box1[2] - box1[0]) * (box1[3] - box1[1])80 area2 = (box2[2] - box2[0]) * (box2[3] - box2[1])81 union = area1 + area2 - intersection82 83 return intersection / union if union > 0 else 084 85 # Apply NMS and keep only highest confidence detection86 for i in range(len(all_detections)):87 if not keep[i]:88 continue89 90 current_box = all_detections[i]['bbox']91 current_label = all_detections[i]['original_label']92 93 for j in range(i + 1, len(all_detections)):94 if not keep[j]:95 continue96 97 # Check if same label type and high IoU98 if (all_detections[j]['original_label'] == current_label and 99 calculate_iou(current_box, all_detections[j]['bbox']) > 0.5):100 # Since list is sorted by confidence, i will always have higher confidence than j101 keep[j] = False102 logging.info(f"Removing duplicate detection of {current_label} with lower confidence "103 f"({all_detections[j]['confidence']:.2f} < {all_detections[i]['confidence']:.2f})")104 105 # Add kept detections to final list106 merged_detections = [det for i, det in enumerate(all_detections) if keep[i]]107 return merged_detections108 109def calculate_font_scale(image_width: int, bbox_width: int) -> float:110 """111 Calculate appropriate font scale based on image and bbox dimensions.112 """113 base_scale = 0.7 # Increased base scale for better visibility114 115 # Adjust font size based on image width and bbox width116 width_ratio = image_width / MAX_DIMENSION117 bbox_ratio = bbox_width / image_width118 119 # Calculate adaptive scale with increased multipliers120 adaptive_scale = base_scale * max(width_ratio, 0.5) * max(bbox_ratio * 6, 0.7)121 122 # Ensure font scale stays within reasonable bounds123 return min(max(adaptive_scale, MIN_FONT_SCALE), MAX_FONT_SCALE)124 125def check_overlap(rect1, rect2):126 """Check if two rectangles overlap."""127 x1_1, y1_1, x2_1, y2_1 = rect1128 x1_2, y1_2, x2_2, y2_2 = rect2129 130 return not (x2_1 < x1_2 or x1_1 > x2_2 or y2_1 < y1_2 or y1_1 > y2_2)131 132def draw_annotation(133 image: np.ndarray,134 bbox: List[int],135 text: str,136 confidence: float,137 model_source: str,138 existing_annotations: List[tuple] = None139) -> None:140 """141 Draw annotation with no background and thin fonts.142 """143 if existing_annotations is None:144 existing_annotations = []145 146 x1, y1, x2, y2 = bbox147 bbox_width = x2 - x1148 image_width = image.shape[1]149 image_height = image.shape[0]150 151 # Calculate adaptive font scale152 font_scale = calculate_font_scale(image_width, bbox_width)153 154 # Simplify the annotation text155 annotation_text = f'{text}\n{confidence:.0f}%'156 lines = annotation_text.split('\n')157 158 # Calculate text dimensions159 font = cv2.FONT_HERSHEY_SIMPLEX160 max_width = 0161 total_height = 0162 line_heights = []163 164 for line in lines:165 (width, height), baseline = cv2.getTextSize(166 line, font, font_scale, THICKNESS167 )168 max_width = max(max_width, width)169 line_height = height + baseline + TEXT_PADDING170 line_heights.append(line_height)171 total_height += line_height172 173 # Calculate initial text position with increased padding174 padding = TEXT_PADDING175 rect_x1 = max(0, x1 - padding)176 rect_x2 = min(image_width, x1 + max_width + padding * 2)177 178 # Try different positions to avoid overlap179 positions = [180 ('top', y1 - total_height - padding),181 ('bottom', y2 + padding),182 ('top_shifted', y1 - total_height - padding * 2),183 ('bottom_shifted', y2 + padding * 2)184 ]185 186 final_position = None187 for pos_name, y_pos in positions:188 if y_pos < 0 or y_pos + total_height > image_height:189 continue190 191 rect = (rect_x1, y_pos, rect_x2, y_pos + total_height)192 overlap = False193 194 for existing_rect in existing_annotations:195 if check_overlap(rect, existing_rect):196 overlap = True197 break198 199 if not overlap:200 final_position = (pos_name, y_pos)201 existing_annotations.append(rect)202 break203 204 # If no non-overlapping position found, use side position205 if final_position is None:206 rect_x1 = max(0, x1 + bbox_width + padding)207 rect_x2 = min(image_width, rect_x1 + max_width + padding * 2)208 y_pos = y1209 final_position = ('side', y_pos)210 211 rect_y1 = final_position[1]212 213 # Draw bounding box (no transparency)214 cv2.rectangle(image, (x1, y1), (x2, y2), BOX_COLOR, BOX_THICKNESS)215 216 # Draw text directly without background217 text_y = rect_y1 + line_heights[0] - padding218 for i, line in enumerate(lines):219 # Draw text with thin lines220 cv2.putText(221 image,222 line,223 (rect_x1 + padding, text_y + sum(line_heights[:i])),224 font,225 font_scale,226 TEXT_COLOR,227 THICKNESS,228 cv2.LINE_AA229 )230 231def run_detection_with_optimal_threshold(232 image_path: str,233 results_dir: str = "results",234 file_name: str = "",235 apply_preprocessing: bool = False,236 resize_image: bool = True, # Changed default to True237 storage: StorageInterface = None238) -> Tuple[str, str, str, List[int]]:239 """Run detection with multiple models and merge results."""240 try:241 image_data = storage.load_file(image_path)242 nparr = np.frombuffer(image_data, np.uint8)243 original_image_cv = cv2.imdecode(nparr, cv2.IMREAD_COLOR)244 image_cv = original_image_cv.copy()245 246 if resize_image:247 logging.info("Resizing image for detection with aspect ratio...")248 image_cv, resized_width, resized_height = resize_image_with_aspect_ratio(image_cv, MAX_DIMENSION)249 else:250 logging.info("Skipping image resizing...")251 resized_height, resized_width = original_image_cv.shape[:2]252 253 if apply_preprocessing:254 logging.info("Preprocessing image for symbol detection...")255 image_cv = preprocess_image_for_symbol_detection(image_cv)256 else:257 logging.info("Skipping image preprocessing for symbol detection...")258 259 all_detections = []260 261 # Run detection with each model262 for model_name, model_path in MODEL_PATHS.items():263 logging.info(f"Running detection with model: {model_name}")264 265 if not model_path:266 logging.warning(f"No model path found for {model_name}")267 continue268 269 model = YOLO(model_path)270 271 best_confidence_threshold = 0.5272 best_detections_list = []273 best_metric = -1274 275 for confidence_threshold in CONFIDENCE_THRESHOLDS:276 logging.info(f"Running detection with confidence threshold: {confidence_threshold}...")277 results = model.predict(source=image_cv, imgsz=MAX_DIMENSION)278 279 detections_list = []280 for result in results:281 for box in result.boxes:282 confidence = float(box.conf[0])283 if confidence >= confidence_threshold:284 x1, y1, x2, y2 = map(float, box.xyxy[0])285 class_id = int(box.cls[0])286 label = result.names[class_id]287 288 scale_x = original_image_cv.shape[1] / resized_width289 scale_y = original_image_cv.shape[0] / resized_height290 x1 *= scale_x291 x2 *= scale_x292 y1 *= scale_y293 y2 *= scale_y294 x1, y1, x2, y2 = map(int, [x1, y1, x2, y2])295 296 split_label = label.split('_')297 if len(split_label) >= 3:298 category = split_label[0]299 type_ = split_label[1]300 new_label = '_'.join(split_label[2:])301 elif len(split_label) == 2:302 category = split_label[0]303 type_ = split_label[1]304 new_label = split_label[1]305 elif len(split_label) == 1:306 category = split_label[0]307 type_ = "Unknown"308 new_label = split_label[0]309 else:310 logging.warning(f"Unexpected label format: {label}. Skipping this detection.")311 continue312 313 detection_id = str(uuid.uuid4())314 detection_info = {315 "symbol_id": detection_id,316 "class_id": class_id,317 "original_label": label,318 "category": category,319 "type": type_,320 "label": new_label,321 "confidence": confidence,322 "bbox": [x1, y1, x2, y2],323 "model_source": model_name324 }325 detections_list.append(detection_info)326 327 metric = evaluate_detections(detections_list)328 if metric > best_metric:329 best_metric = metric330 best_confidence_threshold = confidence_threshold331 best_detections_list = detections_list332 333 all_detections.extend(best_detections_list)334 335 # Merge detections from both models336 merged_detections = merge_detections(all_detections)337 logging.info(f"Total detections after merging: {len(merged_detections)}")338 339 # Draw annotations on the image340 existing_annotations = []341 for det in merged_detections:342 draw_annotation(343 original_image_cv,344 det["bbox"],345 det["original_label"],346 det["confidence"] * 100,347 det["model_source"],348 existing_annotations349 )350 351 # Save results352 storage.create_directory(results_dir)353 file_name_without_extension = os.path.splitext(file_name)[0]354 355 # Prepare output JSON356 total_detected_symbols = len(merged_detections)357 class_counts = {}358 for det in merged_detections:359 full_label = det["original_label"]360 class_counts[full_label] = class_counts.get(full_label, 0) + 1361 362 output_json = {363 "total_detected_symbols": total_detected_symbols,364 "details": class_counts,365 "detections": merged_detections366 }367 368 # Save JSON and image369 detection_json_path = os.path.join(370 results_dir, f'{file_name_without_extension}_detected_symbols.json'371 )372 storage.save_file(373 detection_json_path,374 json.dumps(output_json, indent=4).encode('utf-8')375 )376 377 # Save with maximum quality378 detection_image_path = os.path.join(379 results_dir, f'{file_name_without_extension}_detected_symbols.png' # Using PNG for transparency380 )381 382 # Configure image encoding parameters for maximum quality383 encode_params = [384 cv2.IMWRITE_PNG_COMPRESSION, 0 # No compression for PNG385 ]386 387 # Save as high-quality PNG to preserve transparency388 _, img_encoded = cv2.imencode(389 '.png', 390 original_image_cv,391 encode_params392 )393 394 storage.save_file(detection_image_path, img_encoded.tobytes())395 396 # Calculate diagram bbox from merged detections397 diagram_bbox = [398 min([det['bbox'][0] for det in merged_detections], default=0),399 min([det['bbox'][1] for det in merged_detections], default=0),400 max([det['bbox'][2] for det in merged_detections], default=0),401 max([det['bbox'][3] for det in merged_detections], default=0)402 ]403 404 # Scale up image if it's too small405 min_width = 2000 # Minimum width for good visibility406 if original_image_cv.shape[1] < min_width:407 scale_factor = min_width / original_image_cv.shape[1]408 new_width = min_width409 new_height = int(original_image_cv.shape[0] * scale_factor)410 original_image_cv = cv2.resize(411 original_image_cv, 412 (new_width, new_height), 413 interpolation=cv2.INTER_CUBIC414 )415 416 return (417 detection_image_path,418 detection_json_path,419 f"Total detections after merging: {total_detected_symbols}",420 diagram_bbox421 )422 except Exception as e:423 logging.error(f"An error occurred: {e}")424 return "Error during detection", None, None, None425 426if __name__ == "__main__":427 from storage import StorageFactory428 429 uploaded_file_path = "processed_pages/10219-1-DG-BC-00011.01-REV_A_page_1_text.png"430 results_dir = "results"431 apply_symbol_preprocessing = False432 resize_image = True433 434 storage = StorageFactory.get_storage()435 436 (437 detection_image_path,438 detection_json_path,439 detection_log_message,440 diagram_bbox441 ) = run_detection_with_optimal_threshold(442 uploaded_file_path,443 results_dir=results_dir,444 file_name=os.path.basename(uploaded_file_path),445 apply_preprocessing=apply_symbol_preprocessing,446 resize_image=resize_image,447 storage=storage448 )449 450 logging.info("Detection Image Path: %s", detection_image_path)451 logging.info("Detection JSON Path: %s", detection_json_path)452 logging.info("Detection Log Message: %s", detection_log_message)453 logging.info("Diagram BBox: %s", diagram_bbox)454 logging.info("Done!")