Spaces:
Runtime error
Runtime error
File size: 6,197 Bytes
9cdb8b2 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 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 |