utils/handlers.py

from abc import ABC, abstractmethod
import base64
import io
from typing import Optional, Union
from fastapi import UploadFile
from PIL import Image

class FileHandler(ABC):
    def __init__(self, successor: Optional['FileHandler'] = None):
        self._successor = successor

    @abstractmethod
    async def handle(self, input_data: Union[UploadFile, str]) -> Optional[Image.Image]:
        if self._successor:
            return await self._successor.handle(input_data)
        return None

class ImageFileHandler(FileHandler):
    async def handle(self, input_data: Union[UploadFile, str]) -> Optional[Image.Image]:
        # print(id(UploadFile))
        # print(id(input_data.__class__))
        # print(f"ImageFileHandler received: {type(input_data)}")
        # print(f"Module of input_data: {input_data.__class__.__module__}")
        # print(isinstance(input_data, UploadFile))
        if hasattr(input_data, 'read') and hasattr(input_data, 'filename'):
            print("Handling UploadFile")
            try:
                image_file = await input_data.read()
                return Image.open(io.BytesIO(image_file)).convert("RGB")
            except Exception as e:
                print(f"Error processing UploadFile: {e}")
                return None
        return await super().handle(input_data)

class Base64Handler(FileHandler):
    async def handle(self, input_data: Union[UploadFile, str]) -> Optional[Image.Image]:
        print(f"Base64Handler received: {type(input_data)}")
        if isinstance(input_data, str):
            print("Handling Base64 string")
            try:
                decoded_data = base64.b64decode(input_data)
                # Handle Base64 decoded data (e.g., detect face in the decoded image)
                return Image.open(io.BytesIO(decoded_data)).convert("RGB")
            except Exception:
                pass
        return await super().handle(input_data)