Spaces:
Build error
Build error
File size: 3,498 Bytes
9847531 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 | 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
}
|