Update app.py

app.py

@@ -19,16 +19,22 @@ def predict_fn(img):
         out = model(img)
     probabilities = torch.nn.functional.softmax(out[0], dim=0)
     values, indices = torch.topk(probabilities, k=5)
-    return {LABELS[i]: v.item() for i, v in zip(indices, values)}
+
+    # Get the top labels and select only the first part of the label
+    top_labels = [LABELS[i].split(',')[0] for i in indices]
+
+    # Return only the label and its probability
+    return {top_labels[i]: values[i].item() for i in range(5)}
+
 
 # HTML for the header
 header_html = """
 <div style="text-align: center; max-width: 650px; margin: 0 auto;">
     <img src="https://huggingface.co/spaces/itbetyar/gradio-demo/resolve/main/imgclass.webp" alt="Header Image" style="max-width: 100%; height: auto; margin: 20px 0;">
-    <h1 style="color: #424C5B; font-size: 2.5em;">IT Betyár | Resnet Image Classifier</h1>
+    <h1 style="color: #67768c; font-size: 2.5em;">IT Betyár | Resnet 50 Image Classifier</h1>
     <p style="color: #FFFFFF; font-size: 1.2em;">
-        Üdvözöljük képosztályozónkban! Ez az eszköz egy ImageNeten betanított ResNet50 modellt használ
-        a képek osztályozására. Tölts fel egy képet, és megmutatjuk az 5 legjobb predikciót.
+        Üdvözlünk képosztályozónkban! Ez a minta app egy <b>ResNet50</b> A.I. modellt használ
+        a képek osztályozására. Tölts fel egy képet, és megmutatjuk a három legjobb predikciót.
     </p>
 </div>
 """
@@ -38,14 +44,24 @@ with gr.Blocks() as demo:
     
     with gr.Row():
         with gr.Column():
-            input_image = gr.Image(type='pil')
+            input_image = gr.Image(type='pil', label="Tölts fel egy képet...", width=500, height=400, sources=["upload","webcam"])
             
             with gr.Row():
                 clear_btn = gr.Button("Reset")
                 classify_btn = gr.Button("Mehet")
             
-            # Add examples section properly
-            gr.Examples(
+            
+        with gr.Column():
+            gr.HTML("""<div style="background: #27272A; padding:15px; font-size:16px;">
+                    Alább láthatod a kép osztályozás eredményét</div>""")
+            output = gr.Label(num_top_classes=3, label="A kép osztálya:")
+    
+    classify_btn.click(predict_fn, inputs=input_image, outputs=output)
+    clear_btn.click(lambda: [None, None], inputs=None, outputs=[input_image, output])
+
+# Add examples section properly
+    with gr.Row():
+        gr.Examples(
                 examples=[
                     "imgs/lion.jpg",
                     "imgs/car.jpg", 
@@ -56,13 +72,10 @@ with gr.Blocks() as demo:
                     "imgs/alligator.jpg",
                     "imgs/arc.jpg"
                 ],
-                inputs=input_image,
+                inputs=input_image, label="Betölthető minták"
             )
-            
-        with gr.Column():
-            output = gr.Label()
-    
-    classify_btn.click(predict_fn, inputs=input_image, outputs=output)
-    clear_btn.click(lambda: [None, None], inputs=None, outputs=[input_image, output])
+        
+    with gr.Row():    
+        gr.HTML("""<div style="margin:100px;"></div>""")
 
 demo.launch()