| import torch |
| import numpy as np |
| import cv2 |
| from transformers import AutoModelForImageClassification, AutoFeatureExtractor |
|
|
| class FitPredictor: |
| def __init__(self): |
| self.model_name = "facebook/dinov2-base" |
| self.feature_extractor = AutoFeatureExtractor.from_pretrained(self.model_name) |
| self.model = AutoModelForImageClassification.from_pretrained(self.model_name) |
| self.model.eval() |
| |
| self.fit_types = ['tight', 'regular', 'loose', 'oversized'] |
| |
| self.fit_recommendations = { |
| 'tight': ['Size up for comfort', 'Choose stretchy fabrics', 'Consider relaxed fit alternative'], |
| 'regular': ['True to size', 'Standard fit works for most body types'], |
| 'loose': ['Size down for fitted look', 'Great for layering', 'Can be styled with belt'], |
| 'oversized': ['Intentional oversized style', 'Style with fitted pieces', 'Works well for streetwear'] |
| } |
| |
| self.size_adjustments = { |
| 'tight': '+1', |
| 'regular': '0', |
| 'loose': '-1', |
| 'oversized': '-2' |
| } |
| |
| def predict(self, image: np.ndarray) -> dict: |
| inputs = self.feature_extractor(images=image, return_tensors="pt") |
| |
| with torch.no_grad(): |
| outputs = self.model(**inputs) |
| logits = outputs.logits |
| probabilities = torch.softmax(logits, dim=-1) |
| |
| probs = probabilities[0].tolist() |
| |
| fit_type = self.fit_types[0] |
| confidence = 0.5 |
| for i, ft in enumerate(self.fit_types): |
| if i < len(probs): |
| if probs[i] > confidence: |
| confidence = probs[i] |
| fit_type = ft |
| |
| recommendations = self.fit_recommendations.get(fit_type, self.fit_recommendations['regular']) |
| size_adjustment = self.size_adjustments.get(fit_type, '0') |
| |
| return { |
| "fit": fit_type, |
| "confidence": confidence, |
| "recommendations": recommendations, |
| "size_adjustment": size_adjustment, |
| "body_compatibility": self._get_body_compatibility(fit_type) |
| } |
| |
| def _get_body_compatibility(self, fit_type: str) -> dict: |
| compatibility = { |
| 'tight': {'slim': 0.9, 'athletic': 0.8, 'curvy': 0.6, 'plus': 0.4}, |
| 'regular': {'slim': 0.8, 'athletic': 0.9, 'curvy': 0.8, 'plus': 0.7}, |
| 'loose': {'slim': 0.7, 'athletic': 0.8, 'curvy': 0.9, 'plus': 0.9}, |
| 'oversized': {'slim': 0.9, 'athletic': 0.7, 'curvy': 0.6, 'plus': 0.5} |
| } |
| return compatibility.get(fit_type, compatibility['regular']) |
