scripts/poketwo-classifier.py

import base64
import io
import json
import os
import sys
import threading
import urllib.request
import urllib.parse
from concurrent.futures import ThreadPoolExecutor

os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
os.environ.setdefault("TRANSFORMERS_NO_ADVISORY_WARNINGS", "1")


def emit(payload):
    print(json.dumps(payload, ensure_ascii=True), flush=True)


def load_image(source, image_module):
    if source.startswith("data:"):
        _, encoded = source.split(",", 1)
        raw = base64.b64decode(encoded, validate=True)
    else:
        request = urllib.request.Request(source, headers={"User-Agent": "MU-Bot-Poketwo/1.0"})
        with urllib.request.urlopen(request, timeout=30) as response:
            raw = response.read(16 * 1024 * 1024)
    return image_module.open(io.BytesIO(raw)).convert("RGB")


class SpawnReferenceMatcher:
    def __init__(self):
        try:
            import cv2
            import numpy as np
        except Exception:
            self.cv2 = None
            self.np = None
            return
        self.cv2 = cv2
        self.np = np
        self.detector = cv2.SIFT_create(nfeatures=1200)
        self.matcher = cv2.BFMatcher()
        self.detector_lock = threading.Lock()
        self.reference_cache_lock = threading.Lock()
        self.reference_cache = {}
        self.dataset = os.environ.get(
            "POKETWO_REFERENCE_DATASET",
            "SpreadSheets/Poketwo-Spawn-Images",
        )
        self.candidate_workers = self._env_int("POKETWO_REFERENCE_CANDIDATE_WORKERS", 5, 1, 10)
        self.image_workers = self._env_int("POKETWO_REFERENCE_IMAGE_WORKERS", 2, 1, 4)

    @staticmethod
    def _env_int(name, fallback, minimum, maximum):
        try:
            value = int(os.environ.get(name, ""))
        except ValueError:
            value = fallback
        return max(minimum, min(maximum, value))

    def descriptors(self, image):
        if self.cv2 is None:
            return None
        rgb = self.np.asarray(image.convert("RGB"))
        height, width = rgb.shape[:2]
        crop = rgb[
            int(height * 0.08):max(int(height * 0.92), 1),
            int(width * 0.14):max(int(width * 0.86), 1),
        ]
        crop_height, crop_width = crop.shape[:2]
        if crop_height < 16 or crop_width < 16:
            return None
        mask = self.np.zeros((crop_height, crop_width), self.np.uint8)
        background_model = self.np.zeros((1, 65), self.np.float64)
        foreground_model = self.np.zeros((1, 65), self.np.float64)
        rectangle = (
            max(1, int(crop_width * 0.04)),
            max(1, int(crop_height * 0.04)),
            max(2, int(crop_width * 0.92)),
            max(2, int(crop_height * 0.92)),
        )
        try:
            self.cv2.grabCut(
                crop,
                mask,
                rectangle,
                background_model,
                foreground_model,
                4,
                self.cv2.GC_INIT_WITH_RECT,
            )
            foreground = self.np.where(
                (mask == self.cv2.GC_FGD) | (mask == self.cv2.GC_PR_FGD),
                255,
                0,
            ).astype("uint8")
            foreground = self.cv2.morphologyEx(
                foreground,
                self.cv2.MORPH_CLOSE,
                self.np.ones((5, 5), self.np.uint8),
            )
        except Exception:
            foreground = None
        gray = self.cv2.cvtColor(crop, self.cv2.COLOR_RGB2GRAY)
        # OpenCV does not guarantee that one detector instance is safe across threads.
        with self.detector_lock:
            _, values = self.detector.detectAndCompute(gray, foreground)
        return values

    def _load_reference_descriptors(self, image_url, image_module):
        try:
            reference = load_image(image_url, image_module)
            values = self.descriptors(reference)
            return values if values is not None and len(values) else None
        except Exception:
            return None

    def fetch_reference_descriptors(self, label, image_module):
        cache_key = str(label).strip().lower()
        with self.reference_cache_lock:
            if cache_key in self.reference_cache:
                return self.reference_cache[cache_key]
        folder = urllib.parse.quote(cache_key, safe="")
        dataset = urllib.parse.quote(self.dataset, safe="/")
        api_url = f"https://huggingface.co/api/datasets/{dataset}/tree/main/{folder}?limit=2"
        descriptors = []
        try:
            request = urllib.request.Request(api_url, headers={"User-Agent": "MU-Bot-Poketwo/1.0"})
            with urllib.request.urlopen(request, timeout=20) as response:
                entries = json.loads(response.read().decode("utf-8"))
            image_urls = []
            for entry in entries[:2]:
                path = urllib.parse.quote(str(entry.get("path", "")), safe="/")
                if not path:
                    continue
                image_urls.append(f"https://huggingface.co/datasets/{dataset}/resolve/main/{path}")
            with ThreadPoolExecutor(max_workers=min(self.image_workers, len(image_urls) or 1)) as executor:
                values = executor.map(
                    lambda image_url: self._load_reference_descriptors(image_url, image_module),
                    image_urls,
                )
                descriptors = [value for value in values if value is not None]
        except Exception:
            descriptors = []
        with self.reference_cache_lock:
            return self.reference_cache.setdefault(cache_key, descriptors)

    def score(self, query_descriptors, reference_descriptors):
        if query_descriptors is None or reference_descriptors is None:
            return {"matches": 0, "distance": None}
        try:
            pairs = self.matcher.knnMatch(query_descriptors, reference_descriptors, k=2)
            good = [
                pair[0]
                for pair in pairs
                if len(pair) >= 2 and pair[0].distance < 0.7 * pair[1].distance
            ]
        except Exception:
            return {"matches": 0, "distance": None}
        if not good:
            return {"matches": 0, "distance": None}
        return {
            "matches": len(good),
            "distance": sum(match.distance for match in good) / len(good),
        }

    def rerank(self, image, predictions, image_module):
        query = self.descriptors(image)
        if query is None:
            return predictions
        with ThreadPoolExecutor(max_workers=min(self.candidate_workers, len(predictions) or 1)) as executor:
            references_by_candidate = list(executor.map(
                lambda prediction: self.fetch_reference_descriptors(prediction["label"], image_module),
                predictions,
            ))
        for prediction, references in zip(predictions, references_by_candidate):
            scores = [self.score(query, reference) for reference in references]
            scores = [score for score in scores if score["matches"] > 0]
            if len(scores) >= 2:
                # The sprite is stable across backgrounds; background-only matches are not.
                consistent_matches = min(score["matches"] for score in scores)
                average_distance = sum(score["distance"] for score in scores) / len(scores)
            elif scores:
                consistent_matches = scores[0]["matches"]
                average_distance = scores[0]["distance"]
            else:
                consistent_matches = 0
                average_distance = None
            prediction["visual_reference_count"] = len(scores)
            prediction["visual_matches"] = consistent_matches
            prediction["visual_distance"] = average_distance
        if any(prediction.get("visual_matches", 0) > 0 for prediction in predictions):
            predictions.sort(
                key=lambda item: (
                    item.get("visual_matches", 0),
                    -(item.get("visual_distance") or 9999),
                    item.get("score", 0),
                ),
                reverse=True,
            )
        return predictions


def main():
    try:
        import torch
        from PIL import Image
        from transformers import AutoImageProcessor, AutoModelForImageClassification
    except Exception as error:
        print(f"Missing classifier dependency: {error}", file=sys.stderr, flush=True)
        return 2

    model_id = os.environ.get(
        "POKETWO_SPECIALIST_MODEL",
        "imzynoxprince/pokemons-image-classifier-gen1-gen9",
    )
    try:
        processor = AutoImageProcessor.from_pretrained(model_id)
        model = AutoModelForImageClassification.from_pretrained(model_id)
        model.eval()
    except Exception as error:
        print(f"Could not load classifier {model_id}: {error}", file=sys.stderr, flush=True)
        return 3

    reference_matcher = SpawnReferenceMatcher()
    emit({"type": "ready", "model": model_id, "reference_matching": reference_matcher.cv2 is not None})
    for line in sys.stdin:
        try:
            request = json.loads(line)
            request_id = str(request.get("id", ""))
            image = load_image(str(request.get("image", "")), Image)
            inputs = processor(images=image, return_tensors="pt")
            with torch.inference_mode():
                logits = model(**inputs).logits[0]
                probabilities = torch.softmax(logits, dim=-1)
                count = min(5, int(probabilities.shape[-1]))
                scores, indices = torch.topk(probabilities, k=count)
            predictions = []
            for score, index in zip(scores.tolist(), indices.tolist()):
                label = model.config.id2label.get(index, str(index))
                predictions.append({"label": str(label), "score": float(score)})
            predictions = reference_matcher.rerank(image, predictions, Image)
            emit({"id": request_id, "predictions": predictions})
        except Exception as error:
            emit({"id": str(locals().get("request_id", "")), "error": str(error)})
    return 0


if __name__ == "__main__":
    raise SystemExit(main())