Spaces:
Runtime error
Runtime error
File size: 7,620 Bytes
9cdb8b2 f00d508 9cdb8b2 1139c3b 9cdb8b2 1139c3b 9cdb8b2 1139c3b 9cdb8b2 1139c3b 9cdb8b2 1139c3b 9cdb8b2 1139c3b 9cdb8b2 1139c3b 9cdb8b2 1139c3b 9cdb8b2 1139c3b 9cdb8b2 1139c3b f00d508 9cdb8b2 1139c3b 9cdb8b2 1139c3b 9cdb8b2 1139c3b |
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 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 |
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 CLIPModel, CLIPProcessor
from PIL import Image
class OffTopicDetector:
def __init__(self, model_id: str, device: Optional[str] = None, image_size: str = "E"):
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
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)
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,
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, valid_templates, invalid_classes, autocast)
def predict_probas_item(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)
start = time.time()
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"]]
img_urls = [x.replace("-O.jpg", f"-{self.image_size}.jpg") for x in img_urls]
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
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)
@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
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
|