Update off_topic.py

off_topic.py

@@ -67,13 +67,12 @@ class OffTopicDetector:
         self.image_size = image_size
         self.translator = translator
 
-    def predict_probas(self, images: List[PIL.Image.Image], domain: str,
+    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):
-        site, domain = domain.split("-")
-        domain = re.sub("_", " ", domain).lower()
+        domain = domain.lower()
         if valid_templates:
             valid_classes = [template.format(domain) for template in valid_templates]
         else:
@@ -87,7 +86,7 @@ class OffTopicDetector:
                 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)
+            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"]
         
@@ -130,9 +129,9 @@ class OffTopicDetector:
                             use_title: bool = False,
                 valid_templates: Optional[List[str]] = None,
                 invalid_classes: Optional[List[str]] = None):
-        images, domain, title = self.get_item_data(url_or_id)
+        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, title, valid_templates,
+        probas, valid_probas, invalid_probas = self.predict_probas(images, domain, site, title, valid_templates,
                                                             invalid_classes)
         return images, domain, probas, valid_probas, invalid_probas
 
@@ -146,16 +145,17 @@ class OffTopicDetector:
             item_id = re.sub("-", "", url_or_id)
         start = time.time()
         response = httpx.get(f"https://api.mercadolibre.com/items/{item_id}").json()
-        domain = response["domain_id"]
         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, title
+        return dedup_images, domain_name, site, title
 
     def _filter_dups(self, images: List):
         if len(images) > 1:
@@ -190,24 +190,6 @@ class OffTopicDetector:
             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: