config.py

from dataclasses import dataclass, field
from typing import List, Dict, Optional, Tuple, Union
import numpy as np
from base import BaseConfig


@dataclass
class ImageConfig(BaseConfig):
    """Configuration for global image-related settings"""

    roi: Optional[np.ndarray] = field(default_factory=lambda: np.array([500, 500, 5300, 4000]))
    mask_json_path: str = "./text_and_symbol_bboxes.json"
    save_annotations: bool = True
    annotation_style: Dict = field(default_factory=lambda: {
        'bbox_color': (255, 0, 0),
        'line_color': (0, 255, 0),
        'text_color': (0, 0, 255),
        'thickness': 2,
        'font_scale': 0.6
    })

@dataclass
class SymbolConfig:
    """Configuration for Symbol Detection"""
    model_paths: Dict[str, str] = field(default_factory=lambda: {
        "model1": "models/Intui_SDM_41.pt",
        "model2": "models/Intui_SDM_20.pt"
    })
    # Multiple thresholds to test
    confidence_thresholds: List[float] = field(default_factory=lambda: [0.1, 0.3, 0.5, 0.7, 0.9])
    apply_preprocessing: bool = False
    resize_image: bool = True
    max_dimension: int = 1280
    iou_threshold: float = 0.5
    # Optional mapping from class_id to SymbolType
    symbol_type_mapping: Dict[str, str] = field(default_factory=lambda: {
        "valve": "VALVE",
        "pump": "PUMP",
        "sensor": "SENSOR"
    })


@dataclass
class TagConfig(BaseConfig):
    """Configuration for Tag Detection with OCR"""
    confidence_threshold: float = 0.5
    iou_threshold: float = 0.4
    ocr_engines: List[str] = field(default_factory=lambda: ['tesseract', 'easyocr', 'doctr'])
    text_patterns: Dict[str, str] = field(default_factory=lambda: {
        'Line_Number': r"\d{1,5}-[A-Z]{2,4}-\d{1,3}",
        'Equipment_Tag': r"[A-Z]{1,3}-[A-Z0-9]{1,4}-\d{1,3}",
        'Instrument_Tag': r"\d{2,3}-[A-Z]{2,4}-\d{2,3}",
        'Valve_Number': r"[A-Z]{1,2}-\d{3}",
        'Pipe_Size': r"\d{1,2}\"",
        'Flow_Direction': r"FROM|TO",
        'Service_Description': r"STEAM|WATER|AIR|GAS|DRAIN",
        'Process_Instrument': r"\d{2,3}(?:-[A-Z]{2,3})?-\d{2,3}|[A-Z]{2,3}-\d{2,3}",
        'Nozzle': r"N[0-9]{1,2}|MH",
        'Pipe_Connector': r"[0-9]{1,5}|[A-Z]{1,2}[0-9]{2,5}"
    })
    tesseract_config: str = r'--oem 3 --psm 11'
    easyocr_params: Dict = field(default_factory=lambda: {
        'paragraph': False,
        'height_ths': 2.0,
        'width_ths': 2.0,
        'contrast_ths': 0.2
    })

@dataclass
class LineConfig(BaseConfig):
    """Configuration for Line Detection"""

    threshold_distance: float = 10.0
    expansion_factor: float = 1.1


@dataclass
class PointConfig(BaseConfig):
    """Configuration for Point Detection"""

    threshold_distance: float = 10.0


@dataclass
class JunctionConfig(BaseConfig):
    """Configuration for Junction Detection"""

    window_size: int = 21
    radius: int = 5
    angle_threshold_lb: float = 15.0
    angle_threshold_ub: float = 75.0
    expansion_factor: float = 1.1
    mask_json_paths: List[str] = field(default_factory=list)  # Now a list

    roi: Optional[np.ndarray] = field(default_factory=lambda: np.array([300, 500, 7500, 7000]))
    save_annotations: bool = True
    annotation_style: Dict = None

    def __post_init__(self):
        self.annotation_style = self.annotation_style or {
            'bbox_color': (255, 0, 0),
            'line_color': (0, 255, 0),
            'text_color': (0, 0, 255),
            'thickness': 2,
            'font_scale': 0.6
        }