| import torch |
| import numpy as np |
| from typing import Dict, Tuple |
| import cv2 |
| import utils |
| import model_loader |
| from config import settings |
|
|
|
|
| def predict_character(char_tensor: torch.Tensor, models: Dict[str, torch.nn.Module]) -> str: |
| """Predicts a single character using the Triage and Expert system.""" |
| triage_model = models['triage'] |
|
|
| with torch.no_grad(): |
| triage_output = triage_model(char_tensor) |
| _, predicted_class_idx = torch.max(triage_output, 1) |
| expert_name = settings.TRIAGE_OUTPUT_MAP[predicted_class_idx.item()] |
|
|
| expert_model = models[expert_name] |
| expert_output = expert_model(char_tensor) |
| _, predicted_char_idx = torch.max(expert_output, 1) |
|
|
| offset = settings.EXPERT_LABEL_OFFSETS[expert_name] |
| global_idx = predicted_char_idx.item() + offset |
|
|
| return chr(global_idx) |
|
|
|
|
| def process_image_data(image_data: np.ndarray, models: Dict[str, torch.nn.Module]) -> str: |
| """Performs end-to-end OCR, now with intelligent word spacing.""" |
| gray_image = cv2.cvtColor(image_data, cv2.COLOR_BGR2GRAY) |
| _, binary_image = cv2.threshold(gray_image, 0, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU) |
|
|
| bounding_boxes = utils.segment_characters(binary_image) |
| if not bounding_boxes: |
| return "" |
|
|
| print(f"Found {len(bounding_boxes)} characters to recognize.") |
|
|
| recognized_elements = [] |
| previous_box = bounding_boxes[0] |
|
|
| for i, box in enumerate(bounding_boxes): |
| if i > 0: |
| previous_x, previous_y, previous_w, previous_h = previous_box |
| current_x, current_y, _, _ = box |
|
|
| if current_y > (previous_y + previous_h * settings.NEWLINE_THRESHOLD_FACTOR): |
| recognized_elements.append('\n') |
| elif current_x > (previous_x + previous_w + (previous_w * settings.SPACE_THRESHOLD_FACTOR)): |
| recognized_elements.append(' ') |
|
|
| x, y, w, h = box |
| char_crop = binary_image[y:y + h, x:x + w] |
| char_tensor = utils.prepare_char_for_model(char_crop) |
| predicted_char = predict_character(char_tensor, models) |
| recognized_elements.append(predicted_char) |
|
|
| previous_box = box |
|
|
| return "".join(recognized_elements) |
|
|
