| from PIL import Image |
| import pandas as pd |
| import numpy as np |
| from sentence_transformers import SentenceTransformer |
|
|
| from utils.utils import encode_images, encode_texts |
|
|
| class ClothesClassifier(): |
| def __init__( |
| self, |
| model_name: str, |
| minimum_similarity: float = 0.20, |
| minimum_margin: float = 0.015, |
| ): |
| self.model_name = model_name |
| self.model = SentenceTransformer( |
| self.model_name, |
| device="cpu", |
| ) |
| self.category_prompts = { |
| "dress": [ |
| "a photo of a dress", |
| "a one-piece garment covering the upper and lower body", |
| "a casual or formal dress worn by a person", |
| ], |
| "pants": [ |
| "a photo of pants or trousers", |
| "a lower-body garment with two trouser legs", |
| "jeans, chinos, trousers, or sweatpants", |
| ], |
| "top": [ |
| "a photo of a t-shirt, shirt, blouse, or tank top", |
| "a lightweight garment worn on the upper body", |
| "a shirt that is not a sweater, hoodie, jacket, or coat", |
| ], |
| "sweater_hoodie": [ |
| "a photo of a sweater or hoodie", |
| "a knitted sweater, sweatshirt, or hooded sweatshirt", |
| "a warm pullover garment worn on the upper body", |
| ], |
| "jacket_coat": [ |
| "a photo of a jacket or coat", |
| "an outerwear garment worn over other clothes", |
| "a blazer, jacket, raincoat, or winter coat", |
| ], |
| "shoes": [ |
| "a photo of footwear", |
| "a pair of shoes, sneakers, boots, sandals, or heels", |
| "something worn on the feet", |
| ], |
| "accessories": [ |
| "a photo of a fashion accessory", |
| "a bag, belt, hat, scarf, jewelry, watch, or sunglasses", |
| "an accessory worn with an outfit", |
| ], |
| "skirt": [ |
| "a photo of a skirt", |
| "a lower-body garment that hangs from the waist without separate trouser legs", |
| "a mini skirt, midi skirt, maxi skirt, or pleated skirt", |
| ], |
| } |
| self.labels, self.prototypes = self.build_prototypes() |
| self.minimum_similarity = minimum_similarity |
| self.minimum_margin = minimum_margin |
|
|
|
|
| def build_prototypes( |
| self, |
| ) -> tuple[list[str], np.ndarray]: |
| labels = [] |
| prototypes = [] |
|
|
| for label, prompts in self.category_prompts.items(): |
| prompt_embeddings = encode_texts(self.model, prompts) |
|
|
| prototype = prompt_embeddings.mean(axis=0) |
| prototype /= np.linalg.norm(prototype) |
|
|
| labels.append(label) |
| prototypes.append(prototype) |
|
|
| return labels, np.stack(prototypes) |
|
|
|
|
| def classify( |
| self, |
| images: list[Image.Image] | Image.Image, |
| ) -> pd.DataFrame: |
| if not isinstance(images, list): |
| images = [images] |
| image_embeddings = encode_images(self.model, images) |
| similarities = image_embeddings @ self.prototypes.T |
|
|
| best_indices = similarities.argmax(axis=1) |
| sorted_scores = np.sort(similarities, axis=1) |
|
|
| best_scores = sorted_scores[:, -1] |
| second_best_scores = sorted_scores[:, -2] |
| margins = best_scores - second_best_scores |
|
|
| predictions = [] |
|
|
| for best_index, best_score, margin in zip( |
| best_indices, |
| best_scores, |
| margins, |
| ): |
| if best_score < self.minimum_similarity or margin < self.minimum_margin: |
| predictions.append("other") |
| else: |
| predictions.append(self.labels[best_index]) |
|
|
| result = pd.DataFrame( |
| similarities, |
| columns=self.labels, |
| ) |
|
|
| result["prediction"] = predictions |
| result["best_similarity"] = best_scores |
| result["margin"] = margins |
|
|
| return result |
