DRAPEa / utils /texture_analyzer.py
andevs's picture
Upload 18 files
9c1053c verified
Raw
History Blame Contribute Delete
3.26 kB
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)
}