| """Is this Simba? ConvNeXt-Tiny features (timm) + logistic-regression head.""" | |
| import sys | |
| import joblib, timm, torch | |
| from PIL import Image, ImageOps | |
| from huggingface_hub import hf_hub_download | |
| head = joblib.load(hf_hub_download("makactuary/simba-or-not", "head.joblib")) | |
| backbone = timm.create_model("convnext_tiny.fb_in22k_ft_in1k", pretrained=True, num_classes=0).eval() | |
| transform = timm.data.create_transform(**timm.data.resolve_data_config({}, model=backbone)) | |
| def is_simba(path): | |
| with torch.inference_mode(): | |
| x = backbone(transform(ImageOps.exif_transpose(Image.open(path)).convert("RGB")).unsqueeze(0)).numpy() | |
| p = float(head.predict_proba(x)[0, 1]) | |
| return ("Simba" if p >= 0.5 else "not Simba"), round(p, 3) | |
| if __name__ == "__main__": | |
| print(is_simba(sys.argv[1])) | |