1import onnxruntime as ort
2import onnx
3import numpy as np
4from PIL import Image
5import json
6from huggingface_hub import hf_hub_download
7
8# Download model from HuggingFace (cached automatically)
9MODEL_PATH = hf_hub_download(
10 repo_id="AugustLabs/Author_ID",
11 filename=" style_predictor_500.onnx"
12)
13
14class AuthorID:
15 """
16 Author_ID: Anime Artist Style Recognition
17 Single ONNX file contains: model + centroids + author names
18 """
19
20 def __init__(self, onnx_path):
21 # Load metadata (author names embedded in ONNX)
22 model_onnx = onnx.load(onnx_path)
23 self.names = []
24 self.input_size = 384
25
26 for prop in model_onnx.metadata_props:
27 if prop.key == "author_names":
28 self.names = json.loads(prop.value)
29 elif prop.key == "input_size":
30 self.input_size = int(prop.value)
31
32 providers = ['CUDAExecutionProvider', 'CPUExecutionProvider']
33 self.session = ort.InferenceSession(onnx_path, providers=providers)
34
35 self.mean = np.array([0.485, 0.456, 0.406], dtype=np.float32).reshape(1, 3, 1, 1)
36 self.std = np.array([0.229, 0.224, 0.225], dtype=np.float32).reshape(1, 3, 1, 1)
37
38 def preprocess(self, image_path):
39 img = Image.open(image_path)
40
41 # Handle transparency
42 if img.mode in ('RGBA', 'LA') or (img.mode == 'P' and 'transparency' in img.info):
43 bg = Image.new('RGB', img.size, (255, 255, 255))
44 img = img.convert('RGBA')
45 bg.paste(img, mask=img.split()[3])
46 img = bg
47 else:
48 img = img.convert('RGB')
49
50 img = img.resize((self.input_size, self.input_size), Image.BILINEAR)
51
52 img_np = np.array(img, dtype=np.float32) / 255.0
53 img_np = img_np.transpose(2, 0, 1)[np.newaxis, ...]
54 img_np = (img_np - self.mean) / self.std
55
56 return img_np
57
58 def predict(self, image_path, top_k=5):
59 """Returns list of (author_name, similarity_score)"""
60 img_np = self.preprocess(image_path)
61 top_indices, top_scores = self.session.run(None, {'image': img_np})
62
63 results = []
64 for idx, score in zip(top_indices[0][:top_k], top_scores[0][:top_k]):
65 results.append((self.names[idx], float(score)))
66
67 return results
68
69 def predict_tags(self, image_path, top_k=5):
70 """Returns formatted tags: (artist:name:score)"""
71 results = self.predict(image_path, top_k)
72 return [f"(artist:{name}:{score:.2f})" for name, score in results]
73
74
75# === Example Usage ===
76if __name__ == "__main__":
77 # Initialize (once) — model downloads automatically
78 model = AuthorID(MODEL_PATH)
79
80 # Predict
81 results = model.predict("your_image.jpg", top_k=5)
82
83 print("🎨 Detected artist styles:")
84 for author, score in results:
85 print(f" {author}: {score:.1%}")
86
87 # Or get formatted tags
88 tags = model.predict_tags("your_image.jpg")
89 print("\n📝 Tags:", ", ".join(tags))
