| import torch |
| import torch.nn.functional as F |
| from transformers import ViTForImageClassification, ViTImageProcessor |
| import numpy as np |
| from PIL import Image |
| import json |
|
|
| class ClothingClassifier: |
| def __init__(self): |
| self.model_name = "google/vit-base-patch16-224" |
| self.processor = ViTImageProcessor.from_pretrained(self.model_name) |
| self.model = ViTForImageClassification.from_pretrained(self.model_name) |
| self.model.eval() |
| |
| self.clothing_mapping = { |
| 'shirt': ['shirt', 't-shirt', 'blouse', 'top', 'polo'], |
| 'pants': ['pants', 'trousers', 'jeans', 'slacks', 'chinos'], |
| 'dress': ['dress', 'gown', 'frock', 'sundress'], |
| 'jacket': ['jacket', 'coat', 'blazer', 'hoodie', 'cardigan'], |
| 'shoes': ['shoe', 'sneakers', 'boots', 'sandals', 'loafers'], |
| 'hat': ['hat', 'cap', 'beanie', 'beret', 'fedora'], |
| 'skirt': ['skirt', 'mini', 'midi', 'maxi'], |
| 'shorts': ['shorts', 'bermuda', 'cutoffs'], |
| 'sweater': ['sweater', 'jumper', 'pullover', 'knitwear'], |
| 'accessory': ['bag', 'belt', 'scarf', 'tie', 'gloves'] |
| } |
| |
| self.all_labels = [] |
| for category, keywords in self.clothing_mapping.items(): |
| self.all_labels.extend(keywords) |
| |
| def predict(self, image: np.ndarray) -> dict: |
| inputs = self.processor(images=image, return_tensors="pt") |
| |
| with torch.no_grad(): |
| outputs = self.model(**inputs) |
| logits = outputs.logits |
| probabilities = F.softmax(logits, dim=-1) |
| |
| probs = probabilities[0].tolist() |
| top_k = 10 |
| top_indices = torch.topk(probabilities[0], k=top_k).indices.tolist() |
| |
| predictions = [] |
| confidence_scores = [] |
| |
| for idx in top_indices: |
| label = self.model.config.id2label[idx].lower() |
| confidence = probs[idx] |
| predictions.append(label) |
| confidence_scores.append(confidence) |
| |
| detected_type = self._map_to_clothing_type(predictions) |
| confidence = max(confidence_scores) if confidence_scores else 0.0 |
| |
| return { |
| "type": detected_type, |
| "confidence": confidence, |
| "predictions": predictions[:5], |
| "confidence_scores": confidence_scores[:5] |
| } |
| |
| def _map_to_clothing_type(self, predictions: list) -> str: |
| for pred in predictions: |
| for category, keywords in self.clothing_mapping.items(): |
| if any(keyword in pred.lower() for keyword in keywords): |
| return category |
| return "clothing" |
| |
| def batch_predict(self, images: list) -> list: |
| results = [] |
| for img in images: |
| results.append(self.predict(img)) |
| return results |
