msintui/Intelligent_PID

1import torch2 3from abc import ABC, abstractmethod4from typing import List, Optional, Dict5 6import numpy as np7import cv28 9from pathlib import Path10from loguru import logger11import json12 13from common import DetectionResult14from storage import StorageInterface15from utils import DebugHandler, CoordinateTransformer16 17 18class BaseConfig(ABC):19    """Abstract Base Config for all configuration classes."""20 21    def __post_init__(self):22        """Ensures default values are set correctly for all configs."""23        pass24 25class BaseDetector(ABC):26    """Abstract base class for detection models."""27 28    def __init__(self,29                 config: BaseConfig,30                 debug_handler: DebugHandler = None):31        self.config = config32        self.debug_handler = debug_handler or DebugHandler()33 34    @abstractmethod35    def _load_model(self, model_path: str):36        """Load and return the detection model."""37        pass38 39    @abstractmethod40    def detect(self, image: np.ndarray, *args, **kwargs):41        """Run detection on an input image."""42        pass43 44    @abstractmethod45    def _preprocess(self, image: np.ndarray) -> np.ndarray:46        """Preprocess the input image before detection."""47        pass48 49    @abstractmethod50    def _postprocess(self, image: np.ndarray) -> np.ndarray:51        """Postprocess the input image before detection."""52        pass53 54 55class BaseDetectionPipeline(ABC):56    """Abstract base class for detection pipelines."""57 58    def __init__(59            self,60            storage: StorageInterface,61            debug_handler=None62    ):63        # self.detector = detector64        self.storage = storage65        self.debug_handler = debug_handler or DebugHandler()66        self.transformer = CoordinateTransformer()67 68    @abstractmethod69    def process_image(70            self,71            image_path: str,72            output_dir: str,73            config74    ) -> DetectionResult:75        """Main processing pipeline for a single image."""76        pass77 78    def _apply_roi(self, image: np.ndarray, roi: np.ndarray) -> np.ndarray:79        """Apply region of interest cropping."""80        if roi is not None and len(roi) == 4:81            x_min, y_min, x_max, y_max = roi82            return image[y_min:y_max, x_min:x_max]83        return image84 85    def _adjust_coordinates(self, detections: List[Dict], roi: np.ndarray) -> List[Dict]:86        """Adjust detection coordinates based on ROI"""87        if roi is None or len(roi) != 4:88            return detections89 90        x_offset, y_offset = roi[0], roi[1]91        adjusted = []92 93        for det in detections:94            try:95                adjusted_bbox = [96                    int(det["bbox"][0] + x_offset),97                    int(det["bbox"][1] + y_offset),98                    int(det["bbox"][2] + x_offset),99                    int(det["bbox"][3] + y_offset)100                ]101                adjusted_det = {**det, "bbox": adjusted_bbox}102                adjusted.append(adjusted_det)103            except KeyError:104                logger.warning("Invalid detection format during coordinate adjustment")105        return adjusted106 107    def _persist_results(108            self,109            output_dir: str,110            image_path: str,111            detections: List[Dict],112            annotated_image: Optional[np.ndarray]113    ) -> Dict[str, str]:114        """Save detection results and annotations"""115        self.storage.create_directory(output_dir)116        base_name = Path(image_path).stem117 118        # Save JSON results119        json_path = Path(output_dir) / f"{base_name}_lines.json"120        self.storage.save_file(121            str(json_path),122            json.dumps({123                "solid_lines": {"lines": detections},124                "dashed_lines": {"lines": []}125            }, indent=2).encode('utf-8')126        )127 128        # Save annotated image129        img_path = None130        if annotated_image is not None:131            img_path = Path(output_dir) / f"{base_name}_annotated.jpg"132            _, img_data = cv2.imencode('.jpg', annotated_image)133            self.storage.save_file(str(img_path), img_data.tobytes())134 135        return {136            "json_path": str(json_path),137            "image_path": str(img_path) if img_path else None138        }139