Spaces:
Sleeping
Sleeping
File size: 2,024 Bytes
9cbb8bc |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 |
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()
|