from typing import List

from transformers import pipeline

from internals.util.commons import download_image


class ImageClassifier:
    def __init__(self, candidates: List[str] = ["realistic", "anime", "comic"]):
        self.__candidates = candidates

    def load(self):
        self.pipe = pipeline(
            "zero-shot-image-classification",
            model="philschmid/clip-zero-shot-image-classification",
        )

    def classify(self, image_url: str, width: int, height: int) -> str:
        image = download_image(image_url).resize((width, height))
        results = self.pipe.__call__([image], candidate_labels=self.__candidates)
        results = results[0]
        if len(results) > 0:
            return results[0]["label"]
        return ""