from PIL import Image
import numpy as np
from deepfloorplan_inference import DeepFloorPlanModel

class EndpointModel:
    def __init__(self):
        self.model = DeepFloorPlanModel(model_dir='pretrained')

    def __call__(self, image):
        # image: PIL Image or numpy array
        if isinstance(image, np.ndarray):
            image = Image.fromarray(image)
        result = self.model.predict(image)
        return Image.fromarray(result.astype(np.uint8))

# For Hugging Face Inference Endpoints
model = EndpointModel()

def predict(image):
    return model(image)