Multi image classification

app.py

@@ -126,12 +126,10 @@ model = AutoModelForImageClassification.from_pretrained(MODEL_ID)
 model.eval()
 
 # (Optional) use model's own labels if present
-ID2LABEL = (
-    [model.config.id2label[str(i)] for i in range(model.config.num_labels)]
-    if getattr(model.config, "id2label", None)
-    else ['Eel','Scallop','Crab','Flatfish','Roundfish','Skate','Whelk']
-)
-
+ID2LABEL = [
+    model.config.id2label.get(str(i), model.config.id2label.get(i, f"Label_{i}"))
+    for i in range(model.config.num_labels)
+]
 def classify_image(image):
     if not isinstance(image, Image.Image):
         image = Image.fromarray(image).convert("RGB")
@@ -217,7 +215,7 @@ batch = gr.Interface(
     fn=classify_images_batch,
     inputs=gr.Files(label="Upload up to 10 images"),
     outputs=[
-        gr.Gallery(label="Results (Top-1 in caption)").style(grid=3, height=500),
+        gr.Gallery(label="Results (Top-1 in caption)", height=500, rows=3),
         gr.Dataframe(
             headers=["filename", "top1_label", "top1_conf", "top3_labels", "top3_confs"],
             label="Predictions Table",