update for multi image

app.py

@@ -1,112 +1,235 @@
+# import gradio as gr
+# import torch
+# from transformers import AutoImageProcessor, AutoModelForImageClassification
+# from PIL import Image
+# import os
+
+# # Get model path (Windows compatible)
+# BASE_DIR = os.path.dirname(os.path.abspath(__file__))
+# MODEL_ID = "dshi01/convnext-tiny-224-7clss"
+
+# # Try different possible filenames
+# # possible_names = ['ConvNextmodel.pth', 'convnextmodel.pth', 'ConvNext_model.pth']
+# # model_path = None
+
+# # for name in possible_names:
+# #     test_path = os.path.join(BASE_DIR, name)
+# #     if os.path.exists(test_path):
+# #         model_path = test_path
+# #         print(f"✓ Found model: {name}")
+# #         break
+
+# # if model_path is None:
+# #     raise FileNotFoundError(f"Could not find model file. Tried: {possible_names}")
+
+# # Species categories (7 classes)
+# SPECIES_CATEGORIES = [
+#     'Eel',
+#     'Scallop',
+#     'Crab',
+#     'Flatfish',
+#     'Roundfish',
+#     'Skate',
+#     'Whelk'
+# ]
+
+# # Load model
+# print(f"Loading model from: {MODEL_ID}")
+# # model = AutoModelForImageClassification.from_pretrained(
+# #     'facebook/convnext-tiny-224',
+# #     num_labels=7,
+# #     ignore_mismatched_sizes=True
+# # )
+# processor=AutoImageProcessor.from_pretrained('facebook/convnext-tiny-224')
+# model = AutoModelForImageClassification.from_pretrained(MODEL_ID)
+
+# # Load weights
+# # checkpoint = torch.load(model_path, map_location='cpu', weights_only=False)
+# # if isinstance(checkpoint, dict):
+# #     if 'model' in checkpoint:
+# #         checkpoint = checkpoint['model']
+# #     elif 'state_dict' in checkpoint:
+# #         checkpoint = checkpoint['state_dict']
+
+# # model.load_state_dict(checkpoint, strict=False)
+# # model.eval()
+
+# # Load processor
+# # processor = AutoImageProcessor.from_pretrained('facebook/convnext-tiny-224')
+# # print("✓ Model loaded successfully!")
+
+# def classify_image(image):
+#     """
+#     Classify a benthic species image.
+    
+#     Args:
+#         image: PIL Image or numpy array
+    
+#     Returns:
+#         dict: Predictions with species names and confidence scores
+#     """
+#     # Convert to PIL if needed
+#     if not isinstance(image, Image.Image):
+#         image = Image.fromarray(image).convert('RGB')
+    
+#     # Preprocess
+#     inputs = processor(images=image, return_tensors="pt")
+    
+#     # Predict
+#     with torch.no_grad():
+#         outputs = model(**inputs)
+#         logits = outputs.logits
+#         probabilities = torch.nn.functional.softmax(logits, dim=1)
+    
+#     # Create results dictionary for Gradio
+#     results = {}
+#     for idx, prob in enumerate(probabilities[0]):
+#         results[SPECIES_CATEGORIES[idx]] = float(prob)
+    
+#     return results
+
+# # Create Gradio interface
+# demo = gr.Interface(
+#     fn=classify_image,
+#     inputs=gr.Image(type="pil", label="Upload Underwater Image"),
+#     outputs=gr.Label(num_top_classes=7, label="Species Classification"),
+#     title="🌊 BenthicAI - Benthic Species Classifier",
+#     description="Upload an image of a benthic organism to classify it into one of 7 species categories. Built with ConvNeXT transformer model.",
+#     examples=[
+#         [os.path.join("examples", "eel.jpg")],
+#         [os.path.join("examples", "scallop.jpg")],
+#         [os.path.join("examples", "crab.jpg")],
+#     ] if os.path.exists("examples") else None,
+#     theme=gr.themes.Soft(),
+#     allow_flagging="never"
+# )
+
+# if __name__ == "__main__":
+#     demo.launch(
+#         server_name="0.0.0.0",
+#         server_port=7860,
+#         share=True  # Set to True to get a public URL
+#     )
 import gradio as gr
 import torch
+import torch.nn.functional as F
 from transformers import AutoImageProcessor, AutoModelForImageClassification
 from PIL import Image
 import os
 
-# Get model path (Windows compatible)
 BASE_DIR = os.path.dirname(os.path.abspath(__file__))
 MODEL_ID = "dshi01/convnext-tiny-224-7clss"
 
-# Try different possible filenames
-# possible_names = ['ConvNextmodel.pth', 'convnextmodel.pth', 'ConvNext_model.pth']
-# model_path = None
-
-# for name in possible_names:
-#     test_path = os.path.join(BASE_DIR, name)
-#     if os.path.exists(test_path):
-#         model_path = test_path
-#         print(f"✓ Found model: {name}")
-#         break
-
-# if model_path is None:
-#     raise FileNotFoundError(f"Could not find model file. Tried: {possible_names}")
-
-# Species categories (7 classes)
-SPECIES_CATEGORIES = [
-    'Eel',
-    'Scallop',
-    'Crab',
-    'Flatfish',
-    'Roundfish',
-    'Skate',
-    'Whelk'
-]
-
-# Load model
 print(f"Loading model from: {MODEL_ID}")
-# model = AutoModelForImageClassification.from_pretrained(
-#     'facebook/convnext-tiny-224',
-#     num_labels=7,
-#     ignore_mismatched_sizes=True
-# )
-processor=AutoImageProcessor.from_pretrained('facebook/convnext-tiny-224')
+processor = AutoImageProcessor.from_pretrained("facebook/convnext-tiny-224")
 model = AutoModelForImageClassification.from_pretrained(MODEL_ID)
+model.eval()
 
-# Load weights
-# checkpoint = torch.load(model_path, map_location='cpu', weights_only=False)
-# if isinstance(checkpoint, dict):
-#     if 'model' in checkpoint:
-#         checkpoint = checkpoint['model']
-#     elif 'state_dict' in checkpoint:
-#         checkpoint = checkpoint['state_dict']
+# (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']
+)
 
-# model.load_state_dict(checkpoint, strict=False)
-# model.eval()
+def classify_image(image):
+    if not isinstance(image, Image.Image):
+        image = Image.fromarray(image).convert("RGB")
 
-# Load processor
-# processor = AutoImageProcessor.from_pretrained('facebook/convnext-tiny-224')
-# print("✓ Model loaded successfully!")
+    inputs = processor(images=image, return_tensors="pt")
+    with torch.no_grad():
+        logits = model(**inputs).logits
+        probs = F.softmax(logits, dim=1)[0].tolist()
 
-def classify_image(image):
+    return {ID2LABEL[i]: float(p) for i, p in enumerate(probs)}
+
+# ---------- NEW: batch classify up to 10 images ----------
+MAX_BATCH = 10
+
+def classify_images_batch(files):
     """
-    Classify a benthic species image.
-    
-    Args:
-        image: PIL Image or numpy array
-    
+    files: list of gradio UploadedFile (paths) or None
     Returns:
-        dict: Predictions with species names and confidence scores
+      - gallery: list of (image, caption)
+      - table: list of rows for Dataframe
     """
-    # Convert to PIL if needed
-    if not isinstance(image, Image.Image):
-        image = Image.fromarray(image).convert('RGB')
-    
-    # Preprocess
-    inputs = processor(images=image, return_tensors="pt")
-    
-    # Predict
+    if not files:
+        return [], []
+
+    # Keep at most 10
+    files = files[:MAX_BATCH]
+
+    # Load as PIL
+    pil_images, names = [], []
+    for f in files:
+        path = getattr(f, "name", None) or getattr(f, "path", None) or f
+        try:
+            img = Image.open(path).convert("RGB")
+            pil_images.append(img)
+            names.append(os.path.basename(path))
+        except Exception:
+            # Skip unreadable file
+            continue
+
+    if not pil_images:
+        return [], []
+
+    # Batch preprocess + forward
+    inputs = processor(images=pil_images, return_tensors="pt")
     with torch.no_grad():
-        outputs = model(**inputs)
-        logits = outputs.logits
-        probabilities = torch.nn.functional.softmax(logits, dim=1)
-    
-    # Create results dictionary for Gradio
-    results = {}
-    for idx, prob in enumerate(probabilities[0]):
-        results[SPECIES_CATEGORIES[idx]] = float(prob)
-    
-    return results
+        logits = model(**inputs).logits
+        probs = F.softmax(logits, dim=1)
+
+    # Build outputs
+    gallery = []
+    table_rows = []  # [filename, top1_label, top1_conf, top3_labels, top3_confs]
+
+    for idx, (img, fname) in enumerate(zip(pil_images, names)):
+        p = probs[idx].tolist()
+        top_idxs = sorted(range(len(p)), key=lambda i: p[i], reverse=True)[:3]
+        top1 = top_idxs[0]
+        caption = f"{ID2LABEL[top1]} ({p[top1]:.2%})"
+
+        gallery.append((img, f"{fname}\n{caption}"))
 
-# Create Gradio interface
-demo = gr.Interface(
+        top3_labels = [ID2LABEL[i] for i in top_idxs]
+        top3_scores = [round(p[i], 4) for i in top_idxs]
+        table_rows.append([
+            fname,
+            ID2LABEL[top1],
+            round(p[top1], 4),
+            ", ".join(top3_labels),
+            ", ".join(map(str, top3_scores)),
+        ])
+
+    return gallery, table_rows
+
+# ---------- UI ----------
+single = gr.Interface(
     fn=classify_image,
     inputs=gr.Image(type="pil", label="Upload Underwater Image"),
-    outputs=gr.Label(num_top_classes=7, label="Species Classification"),
-    title="🌊 BenthicAI - Benthic Species Classifier",
-    description="Upload an image of a benthic organism to classify it into one of 7 species categories. Built with ConvNeXT transformer model.",
-    examples=[
-        [os.path.join("examples", "eel.jpg")],
-        [os.path.join("examples", "scallop.jpg")],
-        [os.path.join("examples", "crab.jpg")],
-    ] if os.path.exists("examples") else None,
-    theme=gr.themes.Soft(),
-    allow_flagging="never"
+    outputs=gr.Label(num_top_classes=len(ID2LABEL), label="Species Classification"),
+    title="🌊 BenthicAI - Single Image",
+    description="Classify one image into one of 7 benthic species."
+)
+
+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.Dataframe(
+            headers=["filename", "top1_label", "top1_conf", "top3_labels", "top3_confs"],
+            label="Predictions Table",
+            wrap=True
+        )
+    ],
+    title="🌊 BenthicAI - Batch (up to 10)",
+    description="Upload multiple images (max 10). Outputs a gallery with captions and a table of top predictions.",
 )
 
+demo = gr.TabbedInterface([single, batch], ["Single", "Batch"])
+
 if __name__ == "__main__":
-    demo.launch(
-        server_name="0.0.0.0",
-        server_port=7860,
-        share=True  # Set to True to get a public URL
-    )
+    demo.launch(server_name="0.0.0.0", server_port=7860, share=True)
+