inference.py

from classifier import Classifier
from typing import List
from PIL import Image
import logging

logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
)
logger = logging.getLogger(__name__)

classifier = Classifier()


class Inference():
    def __init__(self):
        self.classifier = classifier

    def _prepare_images(self, images: List[Image.Image]) -> List[Image.Image]:
        """
        Prepare PIL images for classification by converting to RGB
        Args:
            images: List of PIL Image objects
        Returns:
            List of RGB PIL Image objects
        """
        prepared_images = []
        for idx, image in enumerate(images):
            try:
                # Convert to RGB to ensure compatibility
                image = image.convert('RGB')
                prepared_images.append(image)
            except Exception as e:
                raise ValueError(f"Error processing image {idx}: {str(e)}")

        return prepared_images
        
    
    def classify_building(self, images: List[Image.Image], saved_image_paths: List[str] = None) -> dict:
        """
        Classify building type from a list of PIL Image objects
        Args:
            images: List of PIL Image objects
            saved_image_paths: List of paths where images were saved to disk (optional)
        Returns:
            Classification response
        """
        logger.info(f"Preparing {len(images)} images for classification")
        if saved_image_paths:
            logger.info(f"Images saved to disk at: {saved_image_paths}")
        
        prepared_images = self._prepare_images(images)
        logger.info(f"Image preparation successful")
        response = self.classifier.get_response(prepared_images, saved_image_paths)
        return response