| import torch |
| import torch.nn.functional as F |
| from transformers import AutoModelForImageClassification, AutoFeatureExtractor |
| import numpy as np |
| import cv2 |
|
|
| class TextureAnalyzer: |
| def __init__(self): |
| self.model_name = "microsoft/resnet-50" |
| self.feature_extractor = AutoFeatureExtractor.from_pretrained(self.model_name) |
| self.model = AutoModelForImageClassification.from_ppretrained(self.model_name) |
| self.model.eval() |
| |
| self.texture_types = [ |
| 'cotton', 'denim', 'leather', 'silk', 'wool', 'linen', |
| 'polyester', 'velvet', 'knit', 'mesh', 'lace', 'satin', |
| 'chiffon', 'corduroy', 'fleece', 'tweed', 'cashmere' |
| ] |
| |
| self.texture_properties = { |
| 'cotton': {'breathable': True, 'stretchy': False, 'warm': False, 'durable': True}, |
| 'denim': {'breathable': False, 'stretchy': False, 'warm': True, 'durable': True}, |
| 'leather': {'breathable': False, 'stretchy': False, 'warm': True, 'durable': True}, |
| 'silk': {'breathable': True, 'stretchy': False, 'warm': False, 'durable': False}, |
| 'wool': {'breathable': True, 'stretchy': True, 'warm': True, 'durable': True}, |
| 'linen': {'breathable': True, 'stretchy': False, 'warm': False, 'durable': True}, |
| 'polyester': {'breathable': False, 'stretchy': True, 'warm': True, 'durable': True}, |
| 'velvet': {'breathable': False, 'stretchy': False, 'warm': True, 'durable': False} |
| } |
| |
| self.fabric_recommendations = { |
| 'cotton': ['Great for everyday wear', 'Breathable for summer'], |
| 'leather': ['Best for outerwear', 'Requires special care'], |
| 'silk': ['Luxurious feel', 'Dry clean recommended'], |
| 'wool': ['Perfect for winter', 'Keep away from moths'] |
| } |
| |
| def analyze(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 = F.softmax(logits, dim=-1) |
| |
| probs = probabilities[0].tolist() |
| top_indices = torch.topk(probabilities[0], k=5).indices.tolist() |
| |
| texture = self.texture_types[top_indices[0]] if top_indices else 'cotton' |
| confidence = probs[top_indices[0]] if top_indices else 0.5 |
| |
| properties = self.texture_properties.get(texture, {'breathable': False, 'stretchy': False, 'warm': False, 'durable': False}) |
| recommendations = self.fabric_recommendations.get(texture, ['Versatile fabric for various occasions']) |
| |
| return { |
| "texture": texture, |
| "confidence": confidence, |
| "properties": properties, |
| "recommendations": recommendations |
| } |
| |
| def extract_texture_features(self, image: np.ndarray) -> dict: |
| gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY) |
| contrast = gray.std() |
| smoothness = 1 - (gray.var() / (gray.max() + 1e-6)) |
| |
| return { |
| "contrast": float(contrast), |
| "smoothness": float(smoothness), |
| "roughness": float(1 - smoothness) |
| } |