Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| import torch | |
| from PIL import Image | |
| from torchvision import transforms | |
| from model_loader import load_model | |
| from index_to_attr import index_to_attr | |
| # Modell laden | |
| model = load_model("model/AttrPredModel_StateDict.pth") | |
| # taskName pro Index extrahieren | |
| def get_task_map(index_to_attr): | |
| task_map = {} | |
| for idx, desc in index_to_attr.items(): | |
| if "(" in desc and ")" in desc: | |
| task = desc.split("(")[-1].split(")")[0] | |
| task_map[idx] = task | |
| return task_map | |
| task_map = get_task_map(index_to_attr) | |
| # Bildverarbeitungspipeline | |
| preprocess = transforms.Compose([ | |
| transforms.Resize((512, 512)), | |
| transforms.ToTensor(), | |
| transforms.Normalize(mean=[0.6765, 0.6347, 0.6207], | |
| std=[0.3284, 0.3371, 0.3379]) | |
| ]) | |
| # Inferenz-Funktion mit Markierung für unsichere Kategorien | |
| def predict(image): | |
| image = image.convert("RGB") | |
| input_tensor = preprocess(image).unsqueeze(0) | |
| with torch.no_grad(): | |
| output = model(input_tensor) | |
| probs = torch.sigmoid(output).squeeze().numpy() | |
| result = {} | |
| threshold = 0.5 | |
| top_per_task = {} | |
| for idx, score in enumerate(probs): | |
| task = task_map.get(idx, "unknown") | |
| if task not in top_per_task or score > top_per_task[task][1]: | |
| top_per_task[task] = (idx, score) | |
| for task, (idx, score) in top_per_task.items(): | |
| label = index_to_attr.get(idx, f"Unknown ({idx})").split(" (")[0] | |
| result[task] = { | |
| "label": label, | |
| "score": round(float(score), 4), | |
| "confidence": "low" if score < threshold else "high" | |
| } | |
| return result | |
| # Gradio Interface – stabil und einfach | |
| demo = gr.Interface( | |
| fn=predict, | |
| inputs=gr.Image(type="pil", label="Upload image"), | |
| outputs="json", | |
| title="Fashion Attribute Predictor (mit Confidence)", | |
| description="Zeigt pro Attributgruppe die wahrscheinlichste Vorhersage + Confidence ('high' / 'low')." | |
| ) | |
| if __name__ == "__main__": | |
| demo.launch() | |