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 import imagehash from transformers import AutoModelForSeq2SeqLM, AutoTokenizer, CLIPModel, CLIPProcessor from PIL import Image class Translator: def __init__(self, model_id: str, device: Optional[str] = None): self.device = device or ("cuda" if torch.cuda.is_available() else "cpu") self.model_id = model_id self.tokenizer = AutoTokenizer.from_pretrained( model_id) self.model = AutoModelForSeq2SeqLM.from_pretrained(model_id).to(self.device) @property def _bos_token_attr(self): if hasattr(self.tokenizer, "get_lang_id"): return self.tokenizer.get_lang_id elif hasattr(self.tokenizer, "lang_code_to_id"): return self.tokenizer.lang_code_to_id else: return @property def _language_code_mapper(self): if "nllb" in self.model_id.lower(): return {"en": "eng_Latn", "es": "spa_Latn", "pt": "por_Latn"} elif "m2m" in self.model_id.lower(): return {"en": "en", "es": "es", "pt": "pt"} else: return {"en": "eng", "es": "spa", "pt": "por"} def translate(self, texts: List[str], src_lang: str, dest_lang: str = "en", max_length: int = 100): self.tokenizer.src_lang = self._language_code_mapper[src_lang] inputs = self.tokenizer(texts, return_tensors="pt").to(self.device) if "opus" in self.model_id.lower(): forced_bos_token_id = None else: forced_bos_token_id = self._bos_token_attr[self._language_code_mapper["en"]] translated_tokens = self.model.generate( **inputs, forced_bos_token_id=forced_bos_token_id, max_length=max_length ) return self.tokenizer.batch_decode(translated_tokens, skip_special_tokens=True) class OffTopicDetector: def __init__(self, model_id: str, device: Optional[str] = None, image_size: str = "E", translator: Optional[Translator] = 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) self.image_size = image_size self.translator = translator def predict_probas(self, images: List[PIL.Image.Image], domain: str, site: str, title: Optional[str] = None, valid_templates: Optional[List[str]] = None, invalid_classes: Optional[List[str]] = None, autocast: bool = True): domain = domain.lower() 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 title: if site == "CBT": translated_title = title else: if site == "MLB": src_lang = "pt" else: src_lang = "es" translated_title = self.translator.translate(title, src_lang=src_lang, dest_lang="en", max_length=100)[0] valid_classes.append(translated_title.lower()) 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) if self.device == "cuda": torch.cuda.synchronize() 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: 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 if self.device == "cuda": torch.cuda.synchronize() end = time.time() duration = end - start print(f"Model time: {round(duration, 2)} s", f"Model time per image: {round(duration/len(images) * 1000, 0)} 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 predict_probas_url(self, img_urls: List[str], domain: str, site:str, title: Optional[str] = None, valid_templates: Optional[List[str]] = None, invalid_classes: Optional[List[str]] = None, autocast: bool = True): images = self.get_images(img_urls) dedup_images = self._filter_dups(images) return self.predict_probas(images, domain, site, title, valid_templates, invalid_classes, autocast) def predict_probas_item(self, url_or_id: str, use_title: bool = False, valid_templates: Optional[List[str]] = None, invalid_classes: Optional[List[str]] = None): images, domain, site, title = self.get_item_data(url_or_id) title = title if use_title else None probas, valid_probas, invalid_probas = self.predict_probas(images, domain, site, title, 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) start = time.time() response = httpx.get(f"https://api.mercadolibre.com/items/{item_id}").json() title = response["title"] site, domain = response["domain_id"].split("-") img_urls = [x["url"] for x in response["pictures"]] img_urls = [x.replace("-O.jpg", f"-{self.image_size}.jpg") for x in img_urls] domain_name = httpx.get(f"https://api.mercadolibre.com/catalog_domains/CBT-{domain}").json()["name"] end = time.time() duration = end - start print(f"Items API time: {round(duration * 1000, 0)} ms") images = self.get_images(img_urls) dedup_images = self._filter_dups(images) return dedup_images, domain_name, site, title def _filter_dups(self, images: List): if len(images) > 1: hashes = {} for img in images: hashes.update({str(imagehash.average_hash(img)): img}) dedup_hashes = list(dict.fromkeys(hashes)) dedup_images = [img for hash, img in hashes.items() if hash in dedup_hashes] else: dedup_images = images if (diff := len(images) - len(dedup_images)) > 0: print(f"Filtered {diff} images out of {len(images)} due to matching hashes.") return dedup_images 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: {round(duration, 2)} s", f"Download time per image: {round(duration/len(urls) * 1000, 0)} 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) 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