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())