wip

app.py

@@ -1,38 +1,18 @@
 import gradio as gr
 from fastai.vision.all import *
 
-title = """
-    <div style="text-align: center; max-width: 500px; margin: 0 auto;">
-        <div
-        style="
-            display: inline-flex;
-            align-items: center;
-            gap: 0.8rem;
-            font-size: 1.75rem;
-            margin-bottom: 10px;
-        "
-        >
-        <h1 style="font-weight: 600; margin-bottom: 7px;">
-            Grizzly/Black/Teddy bear classification
-        </h1>
-        </div>
-        <p style="margin-bottom: 10px;font-size: 94%;font-weight: 100;line-height: 1.5em;">
-        Please upload a Grizzly, Black or Teddy bear image for classification. 
-        </p>
-    </div>
-"""
+title = "Grizzly/Black/Teddy bear classifier."
 
 learn = load_learner('export.pkl')
-
-categories = ('grizzly', 'black', 'teddy')
+labels = learn.dls.vocab
 
 def classify_image(img):
     pred, idx, probs = learn.predict(img)
-    return dict(zip(categories, map(float, probs)))
+    return {labels[i]: float(probs[i]) for i in range(len(labels))}
 
 gr.HTML(title)
 image = gr.inputs.Image(shape=(224, 224))
-label = gr.outputs.Label()
+label = gr.outputs.Label(num_top_classes=3)
 
-iface = gr.Interface(fn=classify_image, inputs=image, outputs=label)
+iface = gr.Interface(fn=classify_image, inputs=image, outputs=label, title=title)
 iface.launch()