models/material.py

MATERIALS = ["fabric", "leather", "metal", "wood", "plastic", "glass", "ceramic", "paper", "rubber"]

def detect_material_clip(image_bytes):
    from transformers import CLIPProcessor, CLIPModel
    model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32")
    processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")
    image = Image.open(io.BytesIO(image_bytes)).convert('RGB')
    inputs = processor(text=MATERIALS, images=image, return_tensors="pt", padding=True)
    with torch.no_grad():
        outputs = model(**inputs)
    logits = outputs.logits_per_image[0]
    probs = torch.softmax(logits, dim=-1)
    best_idx = probs.argmax().item()
    return {
        "primary_material": MATERIALS[best_idx],
        "confidence": probs[best_idx].item(),
        "all_scores": {mat: prob.item() for mat, prob in zip(MATERIALS, probs)}
    }