from PIL import Image
import torch
from torchvision import transforms  # type: ignore
from transformers import AutoModelForImageSegmentation
from typing import Tuple
import logging

logger = logging.getLogger(__name__)

class ImageSegmenter:
    def __init__(self, image_size: Tuple[int, int] = (1024, 1024)):
        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        logger.info(f"Using device: {self.device}")
        
        try:
            self.birefnet = AutoModelForImageSegmentation.from_pretrained(
                "ZhengPeng7/BiRefNet", 
                trust_remote_code=True
            )
            self.birefnet.to(self.device)
            self.birefnet.eval()  # Set model to evaluation mode
        except Exception as e:
            logger.error(f"Error loading model: {str(e)}")
            raise

        self.transform_image = transforms.Compose([
            transforms.Resize(image_size),
            transforms.ToTensor(),
            transforms.Normalize(
                mean=[0.485, 0.456, 0.406],
                std=[0.229, 0.224, 0.225]
            ),
        ])
    
    async def extract_object(self, image: Image.Image) -> Tuple[Image.Image, Image.Image]:
        try:
            # Transform image
            input_images = self.transform_image(image).unsqueeze(0).to(self.device)

            # Prediction
            with torch.no_grad():
                preds = self.birefnet(input_images)[-1].sigmoid().cpu()
            
            pred = preds[0].squeeze()
            pred_pil = transforms.ToPILImage()(pred)
            mask = pred_pil.resize(image.size)
            image.putalpha(mask)
            
            return image, mask
        except Exception as e:
            logger.error(f"Error in extract_object: {str(e)}")
            raise
    
    async def segment(self, image: Image.Image) -> Tuple[Image.Image, Image.Image]:
        """Fixed typo in method name and added type hints"""
        return await self.extract_object(image)