Update off_topic.py

off_topic.py

@@ -0,0 +1,136 @@
+import re
+import time
+import asyncio
+from io import BytesIO
+from typing import List, Optional
+
+import httpx
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+import PIL
+from transformers import CLIPModel, CLIPProcessor
+from PIL import Image
+
+
+class OffTopicDetector:
+    def __init__(self, model_id: str, device: Optional[str] = None):
+        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
+        self.processor = CLIPProcessor.from_pretrained(model_id)
+        self.model = CLIPModel.from_pretrained(model_id).to(self.device)
+
+    def predict_probas(self, images: List[PIL.Image.Image], domain: str,
+                valid_templates: Optional[List[str]] = None,
+                invalid_classes: Optional[List[str]] = None,
+                autocast: bool = True):
+        if valid_templates:
+            valid_classes = [template.format(domain) for template in valid_templates]
+        else:
+            valid_classes = [f"a photo of {domain}", f"brochure with {domain} image", f"instructions for {domain}", f"{domain} diagram"]
+        if not invalid_classes:
+            invalid_classes = ["promotional ad with store information", "promotional text", "google maps screenshot", "business card", "qr code"]
+        n_valid = len(valid_classes)
+        classes = valid_classes + invalid_classes
+        print(f"Valid classes: {valid_classes}", f"Invalid classes: {invalid_classes}", sep="\n")
+        n_classes = len(classes)
+
+        start = time.time()
+        inputs = self.processor(text=classes, images=images, return_tensors="pt", padding=True).to(self.device)
+        if self.device == "cpu" and autocast is True:
+            print("Disabling autocast due to device='cpu'.")
+            autocast = False
+        with torch.autocast(self.device, enabled=autocast):
+            with torch.no_grad():
+                outputs = self.model(**inputs)
+                probas = outputs.logits_per_image.softmax(dim=1).cpu().numpy() # we can take the softmax to get the label probabilities
+        end = time.time()
+        duration = end - start
+        print(f"Device: {self.device}",
+              f"Response time: {duration}s",
+              f"Response time per image: {round(duration/len(images), 2) * 1000}ms",
+              sep="\n")
+        valid_probas = probas[:, 0:n_valid].sum(axis=1, keepdims=True)
+        invalid_probas = probas[:, n_valid:n_classes].sum(axis=1, keepdims=True)
+        return probas, valid_probas, invalid_probas
+
+    def show(self, images: List[PIL.Image.Image], valid_probas: np.ndarray, n_cols: int = 3, title: Optional[str] = None, threshold: Optional[float] = None):
+        if threshold is not None:
+            prediction = self.apply_threshold(valid_probas, threshold)
+            title_scores = [f"Valid: {pred.squeeze()}" for pred in prediction]
+        else:
+            prediction = np.round(valid_probas[:, 0], 2)
+            title_scores = [f"Valid: {pred:.2f}" for pred in prediction]
+        n_images = len(images)
+        n_rows = int(np.ceil(n_images / n_cols))
+        fig, axes = plt.subplots(n_rows, n_cols, figsize=(16, 16))
+        for i, ax in enumerate(axes.ravel()):
+            ax.axis("off")
+            try:
+                ax.imshow(images[i])
+                ax.set_title(title_scores[i])
+            except IndexError:
+                continue
+        if title:
+            fig.suptitle(title)
+        fig.tight_layout()
+        return
+
+    def predict_item_probas(self, url_or_id: str,
+                valid_templates: Optional[List[str]] = None,
+                invalid_classes: Optional[List[str]] = None):
+        images, domain = self.get_item_data(url_or_id)
+        probas, valid_probas, invalid_probas = self.predict_probas(images, domain, valid_templates,
+                                                            invalid_classes)
+        return images, domain, probas, valid_probas, invalid_probas
+
+    def apply_threshold(self, valid_probas: np.ndarray, threshold: float = 0.4):
+        return valid_probas >= threshold
+
+    def get_item_data(self, url_or_id: str):
+        if url_or_id.startswith("http"):
+            item_id = "".join(url_or_id.split("/")[3].split("-")[:2])
+        else:
+            item_id = re.sub("-", "", url_or_id)
+        response = httpx.get(f"https://api.mercadolibre.com/items/{item_id}").json()
+        domain = re.sub("_", " ", response["domain_id"].split("-")[-1]).lower()
+        img_urls = [x["url"] for x in response["pictures"]]
+        images = self.get_images(img_urls)
+        return images, domain
+
+    def get_images(self, urls: List[str]):
+        start = time.time()
+        images = asyncio.run(self._gather_download_tasks(urls))
+        end = time.time()
+        duration = end - start
+        print(f"Download time: {duration}s",
+              f"Download time per image: {round(duration/len(urls), 2) * 1000}ms",
+              sep="\n")
+        return asyncio.run(self._gather_download_tasks(urls))
+
+    async def _gather_download_tasks(self, urls: List[str]):
+
+        async def _process_download(url: str, client: httpx.AsyncClient):
+            response = await client.get(url)
+            return Image.open(BytesIO(response.content))
+
+        async with httpx.AsyncClient() as client:
+            tasks = [_process_download(url, client) for url in urls]
+            return await asyncio.gather(*tasks)
+
+    @staticmethod
+    def _non_async_get_item_data(url_or_id: str, save_images: bool = False):
+        if url_or_id.startswith("http"):
+            item_id = "".join(url_or_id.split("/")[3].split("-")[:2])
+        else:
+            item_id = re.sub("-", "", url_or_id)
+        response = httpx.get(f"https://api.mercadolibre.com/items/{item_id}").json()
+        domain = re.sub("_", " ", response["domain_id"].split("-")[-1]).lower()
+        img_urls = [x["url"] for x in response["pictures"]]
+        images = []
+        for img_url in img_urls:
+            img = httpx.get(img_url)
+            images.append(Image.open(BytesIO(img.content)))
+            if save_images:
+                with open(re.sub("D_NQ_NP_", "", img_url.split("/")[-1]) , "wb") as f:
+                    f.write(img.content)
+        return images, domain