| import logging |
|
|
| log = logging.getLogger("tell") |
|
|
| |
| LABELS = { |
| "chart": "a bar chart or line graph with axes", |
| "screenshot": "a screenshot of a social media post or a text message conversation", |
| "photo": "a photo, painting, meme, or AI-generated picture", |
| } |
|
|
| _clf = None |
|
|
|
|
| def classify_image(image, threshold: float = 0.5) -> tuple[str, float]: |
| global _clf |
| try: |
| if _clf is None: |
| from transformers import pipeline |
| _clf = pipeline("zero-shot-image-classification", model="openai/clip-vit-base-patch32") |
| top = _clf(image, candidate_labels=list(LABELS.values()))[0] |
| except Exception as e: |
| log.warning("Classifier failed: %s", e) |
| return "unknown", 0.0 |
|
|
| if top["score"] < threshold: |
| return "unknown", top["score"] |
| kind = next(k for k, v in LABELS.items() if v == top["label"]) |
| return kind, top["score"] |
