File size: 10,204 Bytes
2b54cc5 80eccf9 2b54cc5 80eccf9 2b54cc5 80eccf9 2b54cc5 80eccf9 2b54cc5 80eccf9 2b54cc5 80eccf9 2b54cc5 80eccf9 2b54cc5 80eccf9 2b54cc5 80eccf9 2b54cc5 80eccf9 2b54cc5 80eccf9 2b54cc5 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 | 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())
|