line_detectors.py

import cv2
import numpy as np
import torch
from typing import Dict, List, Optional, Tuple
from loguru import logger

# Check if DeepLSD is available
try:
    from deeplsd.models.deeplsd_inference import DeepLSD
    DEEPLSD_AVAILABLE = True
except ImportError:
    DEEPLSD_AVAILABLE = False
    logger.warning("DeepLSD not available, falling back to OpenCV")


class OpenCVLineDetector:
    """Fallback line detector using OpenCV's HoughLinesP"""
    
    def __init__(self):
        self.params = {
            'threshold': 50,
            'minLineLength': 50,
            'maxLineGap': 10
        }

    def detect(self, image: np.ndarray) -> Dict:
        """Detect lines using HoughLinesP"""
        if len(image.shape) == 3:
            gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
        else:
            gray = image
            
        edges = cv2.Canny(gray, 50, 150, apertureSize=3)
        lines = cv2.HoughLinesP(
            edges, 1, np.pi/180,
            threshold=self.params['threshold'],
            minLineLength=self.params['minLineLength'],
            maxLineGap=self.params['maxLineGap']
        )
        
        detections = []
        if lines is not None:
            for line in lines:
                x1, y1, x2, y2 = line[0]
                detections.append({
                    'x1': float(x1),
                    'y1': float(y1),
                    'x2': float(x2),
                    'y2': float(y2),
                    'confidence': 1.0
                })
                
        return {'lines': detections}


class DeepLSDDetector:
    """Line detector using DeepLSD model"""
    
    def __init__(self, model_path: str):
        if not DEEPLSD_AVAILABLE:
            raise ImportError("DeepLSD is not available")
            
        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        self.model = self._load_model(model_path)

    def _load_model(self, model_path: str) -> DeepLSD:
        """Load the DeepLSD model"""
        try:
            ckpt = torch.load(model_path, map_location=self.device)
            model = DeepLSD()
            model.load_state_dict(ckpt['model'])
            return model.to(self.device).eval()
        except Exception as e:
            logger.error(f"Failed to load DeepLSD model: {str(e)}")
            raise

    def detect(self, image: np.ndarray) -> Dict:
        """Detect lines using DeepLSD"""
        try:
            # Convert to tensor
            if len(image.shape) == 3:
                gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
            else:
                gray = image
                
            tensor = torch.tensor(gray, dtype=torch.float32, device=self.device)[None, None] / 255.0
            
            # Run inference
            with torch.no_grad():
                output = self.model({"image": tensor})
                lines = output["lines"][0]  # [N, 2, 2] array
                
            # Convert to standard format
            detections = []
            for line in lines:
                (x1, y1), (x2, y2) = line
                detections.append({
                    'x1': float(x1),
                    'y1': float(y1),
                    'x2': float(x2),
                    'y2': float(y2),
                    'confidence': float(output.get("confidence", [1.0])[0])
                })
                
            return {'lines': detections}
            
        except Exception as e:
            logger.error(f"Error in DeepLSD detection: {str(e)}")
            return {'lines': []}