app/services/image_processor.py

import cv2
import numpy as np
import torch
from typing import Tuple
from app.core.config import settings
from app.core.exceptions import ImageProcessingError
import logging

logger = logging.getLogger(__name__)

class ImageProcessor:
    def __init__(self):
        self.target_size = (settings.IMAGE_WIDTH, settings.IMAGE_HEIGHT)
        self.max_size = settings.MAX_IMAGE_SIZE
        self.original_image = None
    
    def validate_image(self, image_bytes: bytes) -> None:
        if len(image_bytes) > self.max_size:
            raise ImageProcessingError(f"Image size exceeds maximum allowed size of {self.max_size} bytes")
    
    def decode_image(self, image_bytes: bytes) -> np.ndarray:
        try:
            nparr = np.frombuffer(image_bytes, np.uint8)
            # First try to read as grayscale for dental X-rays
            image_gray = cv2.imdecode(nparr, cv2.IMREAD_GRAYSCALE)
            if image_gray is None:
                raise ImageProcessingError("Failed to decode image")
            
            # Convert grayscale to 3-channel RGB for model input
            # This ensures proper handling of dental X-rays
            image = cv2.cvtColor(image_gray, cv2.COLOR_GRAY2RGB)
            
            # Store original for overlay
            self.original_image = image.copy()
            
            logger.info(f"Decoded image shape: {image.shape}, dtype: {image.dtype}")
            return image
        except Exception as e:
            raise ImageProcessingError(f"Error decoding image: {str(e)}")
    
    def preprocess(self, image: np.ndarray) -> torch.Tensor:
        try:
            # Apply histogram equalization for better contrast (important for X-rays)
            if len(image.shape) == 3:
                # Convert to YUV and equalize Y channel
                image_yuv = cv2.cvtColor(image, cv2.COLOR_RGB2YUV)
                image_yuv[:,:,0] = cv2.equalizeHist(image_yuv[:,:,0])
                image = cv2.cvtColor(image_yuv, cv2.COLOR_YUV2RGB)
            
            resized = cv2.resize(image, self.target_size, interpolation=cv2.INTER_LANCZOS4)
            normalized = resized.astype(np.float32) / 255.0
            
            # Store resized original for overlay
            self.original_image = resized.copy()
            
            transposed = np.transpose(normalized, (2, 0, 1))
            tensor = torch.from_numpy(transposed).unsqueeze(0)
            logger.info(f"Preprocessed tensor shape: {tensor.shape}")
            return tensor
        except Exception as e:
            raise ImageProcessingError(f"Error preprocessing image: {str(e)}")
    
    def postprocess_mask(self, mask: np.ndarray, overlay: bool = True, alpha: float = 0.6) -> np.ndarray:
        """
        Convert segmentation mask to colored output with optional overlay
        Args:
            mask: 2D segmentation mask
            overlay: Whether to overlay on original image
            alpha: Opacity of overlay (0-1), higher = more mask visible
        """
        h, w = mask.shape
        colored_mask = np.zeros((h, w, 3), dtype=np.uint8)
        
        # Apply colors to each class
        for idx, color in enumerate(settings.COLORMAP):
            colored_mask[mask == idx] = color
        
        # If overlay is requested and we have the original image
        if overlay and self.original_image is not None:
            # Ensure original image is same size
            if self.original_image.shape[:2] != (h, w):
                original_resized = cv2.resize(self.original_image, (w, h))
            else:
                original_resized = self.original_image
            
            # Create overlay: blend original with colored mask
            # For background pixels (class 0), show more of the original
            background_mask = (mask == 0)
            overlay_output = colored_mask.copy()
            
            # Apply alpha blending
            overlay_output = cv2.addWeighted(
                original_resized.astype(np.uint8), 
                1 - alpha, 
                colored_mask, 
                alpha, 
                0
            )
            
            # For foreground classes, increase visibility
            for idx in range(1, len(settings.COLORMAP)):
                class_mask = (mask == idx)
                if np.any(class_mask):
                    # Make segmented regions more visible
                    overlay_output[class_mask] = cv2.addWeighted(
                        original_resized[class_mask].astype(np.uint8),
                        0.3,
                        colored_mask[class_mask],
                        0.7,
                        0
                    )
            
            return overlay_output
        
        return colored_mask
    
    def encode_image(self, image: np.ndarray, format: str = ".jpg") -> bytes:
        try:
            encode_param = [int(cv2.IMWRITE_JPEG_QUALITY), 95]
            if format == ".png":
                encode_param = [int(cv2.IMWRITE_PNG_COMPRESSION), 9]
            
            success, encoded = cv2.imencode(format, image, encode_param)
            if not success:
                raise ImageProcessingError("Failed to encode image")
            return encoded.tobytes()
        except Exception as e:
            raise ImageProcessingError(f"Error encoding image: {str(e)}")

image_processor = ImageProcessor()