Track coral images with Git LFS

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+coral_images/*.jpg filter=lfs diff=lfs merge=lfs -text

app _bk.py

@@ -2,54 +2,104 @@ import gradio as gr
 import numpy as np
 from PIL import Image, ImageDraw
 import torch
-from torchvision import transforms
-from transformers import AutoModelForImageClassification, AutoFeatureExtractor
-
-# Define all available models
-MODEL_LIST = {
-    'beit': "microsoft/beit-base-patch16-224-pt22k-ft22k",
-    'vit': "google/vit-base-patch16-224",
-    'convnext': "facebook/convnext-tiny-224",
-}
-
-# Global variables
-current_model = None
-current_preprocessor = None
-device = "cuda" if torch.cuda.is_available() else "cpu"  # Dynamically set device
-
-# Load model and preprocessor
-def load_model_and_preprocessor(model_name):
-    """Load model and preprocessor for a given model name."""
-    global current_model, current_preprocessor
-    print(f"Loading model and preprocessor for: {model_name} on {device}")
-    current_model = AutoModelForImageClassification.from_pretrained(MODEL_LIST[model_name]).to(device).eval()
-    current_preprocessor = AutoFeatureExtractor.from_pretrained(MODEL_LIST[model_name])
-    return f"Model {model_name} loaded successfully on {device}."
-
-# Predict function
-def predict(image, model, preprocessor):
-    """Make a prediction on the given image patch using the loaded model."""
-    if model is None or preprocessor is None:
-        raise ValueError("Model and preprocessor are not loaded.")
-    inputs = preprocessor(images=image, return_tensors="pt").to(device)
+import torchvision.transforms as transforms
+import timm
+
+# URL for the Hugging Face checkpoint
+CHECKPOINT_URL = "https://huggingface.co/ReefNet/beit_global/resolve/main/checkpoint-60.pth"
+
+# Class labels
+all_classes = [
+    'Acanthastrea', 'Acropora', 'Agaricia', 'Alveopora', 'Astrea', 'Astreopora',
+    'Caulastraea', 'Coeloseris', 'Colpophyllia', 'Coscinaraea', 'Ctenactis',
+    'Cycloseris', 'Cyphastrea', 'Dendrogyra', 'Dichocoenia', 'Diploastrea',
+    'Diploria', 'Dipsastraea', 'Echinophyllia', 'Echinopora', 'Euphyllia',
+    'Eusmilia', 'Favia', 'Favites', 'Fungia', 'Galaxea', 'Gardineroseris',
+    'Goniastrea', 'Goniopora', 'Halomitra', 'Herpolitha', 'Hydnophora',
+    'Isophyllia', 'Isopora', 'Leptastrea', 'Leptoria', 'Leptoseris',
+    'Lithophyllon', 'Lobactis', 'Lobophyllia', 'Madracis', 'Meandrina', 'Merulina',
+    'Montastraea', 'Montipora', 'Mussa', 'Mussismilia', 'Mycedium', 'Orbicella',
+    'Oulastrea', 'Oulophyllia', 'Oxypora', 'Pachyseris', 'Pavona', 'Pectinia',
+    'Physogyra', 'Platygyra', 'Plerogyra', 'Plesiastrea', 'Pocillopora',
+    'Podabacia', 'Porites', 'Psammocora', 'Pseudodiploria', 'Sandalolitha',
+    'Scolymia', 'Seriatopora', 'Siderastrea', 'Stephanocoenia', 'Stylocoeniella',
+    'Stylophora', 'Tubastraea', 'Turbinaria'
+]
+
+# Function to load the BeIT model
+def load_model(model_name):
+    print(f"Loading {model_name} model...")
+    if model_name == 'beit':
+        args = type('', (), {})()
+        args.model = 'beitv2_large_patch16_224.in1k_ft_in22k_in1k'
+        args.nb_classes = len(all_classes)
+        args.drop_path = 0.1
+
+        # Create model
+        model = timm.create_model(
+            args.model,
+            pretrained=False,
+            num_classes=args.nb_classes,
+            drop_path_rate=args.drop_path,
+            use_rel_pos_bias=True,
+            use_abs_pos_emb=True,
+        )
+
+        # Load checkpoint from Hugging Face
+        checkpoint = torch.hub.load_state_dict_from_url(CHECKPOINT_URL, map_location="cpu")
+        state_dict = checkpoint.get('model', checkpoint)
+
+        # Filter state dict
+        filtered_state_dict = {k: v for k, v in state_dict.items() if "relative_position_index" not in k}
+        model.load_state_dict(filtered_state_dict, strict=False)
+    else:
+        raise ValueError(f"Model {model_name} not implemented!")
+
+    # Move model to CUDA if available
+    model.eval()
+    if torch.cuda.is_available():
+        model.cuda()
+    return model
+
+# Preprocessing transforms
+preprocess = transforms.Compose([
+    transforms.Resize((224, 224)),
+    transforms.ToTensor(),
+    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+])
+
+# Initialize selected model
+selected_model_name = 'beit'
+model = load_model(selected_model_name)
+
+def predict_label(image):
+    """Predict the label for the given image."""
+    # Ensure the image is a PIL Image
+    if isinstance(image, np.ndarray):
+        image = Image.fromarray(image)
+    elif not isinstance(image, Image.Image):
+        raise TypeError(f"Unexpected type {type(image)}, expected PIL.Image or numpy.ndarray.")
+
+    input_tensor = preprocess(image).unsqueeze(0)
+    if torch.cuda.is_available():
+        input_tensor = input_tensor.cuda()
+
     with torch.no_grad():
-        outputs = model(**inputs)
-        predicted_class = torch.argmax(outputs.logits, dim=1).item()
-    return model.config.id2label[predicted_class]
+        outputs = model(input_tensor)
+        predicted_class = torch.argmax(outputs, dim=1).item()
+
+    return all_classes[predicted_class]
+
 
 # Function to draw a rectangle on the image
 def draw_rectangle(image, x, y, size=224):
-    """Draw a rectangle on the image."""
-    image_pil = image.copy()  # Create a copy to avoid modifying the original image
+    image_pil = image.copy()
     draw = ImageDraw.Draw(image_pil)
-    x1, y1 = x, y
-    x2, y2 = x + size, y + size
-    draw.rectangle([x1, y1, x2, y2], outline="red", width=5)
+    draw.rectangle([x, y, x + size, y + size], outline="red", width=3)
     return image_pil
 
-# Function to crop the image
+# Crop a region of interest
 def crop_image(image, x, y, size=224):
-    """Crop a region from the image."""
     image_np = np.array(image)
     h, w, _ = image_np.shape
     x = min(max(x, 0), w - size)
@@ -57,55 +107,40 @@ def crop_image(image, x, y, size=224):
     cropped = image_np[y:y+size, x:x+size]
     return Image.fromarray(cropped)
 
-# Gradio Interface
+# Gradio UI
 with gr.Blocks() as demo:
-    gr.Markdown("## Test Public Models for Coral Classification")
-
+    gr.Markdown("## Coral Classification with BeIT Model")
     with gr.Row():
         with gr.Column():
-            model_selector = gr.Dropdown(choices=list(MODEL_LIST.keys()), value='beit', label="Select Model")
             image_input = gr.Image(type="pil", label="Upload Image", interactive=True)
-            x_slider = gr.Slider(minimum=0, maximum=1000, step=1, value=0, label="X Coordinate")
-            y_slider = gr.Slider(minimum=0, maximum=1000, step=1, value=0, label="Y Coordinate")
+            x_slider = gr.Slider(0, 1000, step=1, value=0, label="X Coordinate")
+            y_slider = gr.Slider(0, 1000, step=1, value=0, label="Y Coordinate")
         with gr.Column():
-            interactive_image = gr.Image(label="Interactive Image with Selection")
+            interactive_image = gr.Image(label="Interactive Image")
             cropped_image = gr.Image(label="Cropped Patch")
             label_output = gr.Textbox(label="Predicted Label")
-
-    # Update the model and preprocessor
-    def update_model(model_name):
-        return load_model_and_preprocessor(model_name)
-
-    # Update the rectangle and crop the patch
+    
+    # Interactions
     def update_selection(image, x, y):
         overlay_image = draw_rectangle(image, x, y)
         cropped = crop_image(image, x, y)
         return overlay_image, cropped
 
-    # Predict the label from the cropped patch
     def predict_from_cropped(cropped):
-        print(f"Type of cropped_image before prediction: {type(cropped)}")
-        return predict(cropped, current_model, current_preprocessor)
+        return predict_label(cropped)
 
-    # Buttons and interactions
     crop_button = gr.Button("Crop")
     crop_button.click(fn=update_selection, inputs=[image_input, x_slider, y_slider], outputs=[interactive_image, cropped_image])
 
     predict_button = gr.Button("Predict")
     predict_button.click(fn=predict_from_cropped, inputs=cropped_image, outputs=label_output)
 
-    model_selector.change(fn=update_model, inputs=model_selector, outputs=None)
-
-    # Update sliders dynamically based on uploaded image size
     def update_sliders(image):
-        if image is not None:
+        if image:
             width, height = image.size
             return gr.update(maximum=width - 224), gr.update(maximum=height - 224)
         return gr.update(), gr.update()
 
     image_input.change(fn=update_sliders, inputs=image_input, outputs=[x_slider, y_slider])
 
-    # Initialize model on app start
-    demo.load(fn=lambda: load_model_and_preprocessor('beit'), inputs=None, outputs=None)
-
 demo.launch(server_name="0.0.0.0", server_port=7860)

app.py

@@ -26,34 +26,41 @@ all_classes = [
     'Stylophora', 'Tubastraea', 'Turbinaria'
 ]
 
+# Example image paths
+example_images = {
+    "Acropora": "coral_images/Acropora_millepora.jpg",
+    "Agaricia": "coral_images/Agaricia_agaricites.jpg",
+    "Acropora": "coral_images/Acropora_aculeus.jpg",
+    "Montipora": "coral_images/Montipora_patula.jpg",
+    "Pocillopora": "coral_images/Pocillopora_acuta.jpg",
+    "Porites": "coral_images/porities_lobata.jpg",
+    "Favites": "coral_images/Favites_abdita.jpg",
+    "Fungia": "coral_images/Fungia_concinna.jpg",
+}
+
 # Function to load the BeIT model
 def load_model(model_name):
     print(f"Loading {model_name} model...")
-    if model_name == 'beit':
-        args = type('', (), {})()
-        args.model = 'beitv2_large_patch16_224.in1k_ft_in22k_in1k'
-        args.nb_classes = len(all_classes)
-        args.drop_path = 0.1
-
-        # Create model
-        model = timm.create_model(
-            args.model,
-            pretrained=False,
-            num_classes=args.nb_classes,
-            drop_path_rate=args.drop_path,
-            use_rel_pos_bias=True,
-            use_abs_pos_emb=True,
-        )
-
-        # Load checkpoint from Hugging Face
-        checkpoint = torch.hub.load_state_dict_from_url(CHECKPOINT_URL, map_location="cpu")
-        state_dict = checkpoint.get('model', checkpoint)
-
-        # Filter state dict
-        filtered_state_dict = {k: v for k, v in state_dict.items() if "relative_position_index" not in k}
-        model.load_state_dict(filtered_state_dict, strict=False)
-    else:
-        raise ValueError(f"Model {model_name} not implemented!")
+    args = type('', (), {})()
+    args.model = 'beitv2_large_patch16_224.in1k_ft_in22k_in1k'
+    args.nb_classes = len(all_classes)
+    args.drop_path = 0.1
+
+    # Create model
+    model = timm.create_model(
+        args.model,
+        pretrained=False,
+        num_classes=args.nb_classes,
+        drop_path_rate=args.drop_path,
+        use_rel_pos_bias=True,
+        use_abs_pos_emb=True,
+    )
+
+    # Load checkpoint from Hugging Face
+    checkpoint = torch.hub.load_state_dict_from_url(CHECKPOINT_URL, map_location="cpu")
+    state_dict = checkpoint.get('model', checkpoint)
+    filtered_state_dict = {k: v for k, v in state_dict.items() if "relative_position_index" not in k}
+    model.load_state_dict(filtered_state_dict, strict=False)
 
     # Move model to CUDA if available
     model.eval()
@@ -68,18 +75,12 @@ preprocess = transforms.Compose([
     transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
 ])
 
-# Initialize selected model
-selected_model_name = 'beit'
-model = load_model(selected_model_name)
+# Initialize model
+model = load_model('beit')
 
 def predict_label(image):
-    """Predict the label for the given image."""
-    # Ensure the image is a PIL Image
     if isinstance(image, np.ndarray):
         image = Image.fromarray(image)
-    elif not isinstance(image, Image.Image):
-        raise TypeError(f"Unexpected type {type(image)}, expected PIL.Image or numpy.ndarray.")
-
     input_tensor = preprocess(image).unsqueeze(0)
     if torch.cuda.is_available():
         input_tensor = input_tensor.cuda()
@@ -90,15 +91,16 @@ def predict_label(image):
 
     return all_classes[predicted_class]
 
-
-# Function to draw a rectangle on the image
 def draw_rectangle(image, x, y, size=224):
-    image_pil = image.copy()
+    """Draw a clear red rectangle with increased thickness."""
+    image_pil = image.copy()  # Create a copy to avoid modifying the original image
     draw = ImageDraw.Draw(image_pil)
-    draw.rectangle([x, y, x + size, y + size], outline="red", width=3)
+    x1, y1 = x, y
+    x2, y2 = x + size, y + size
+    draw.rectangle([x1, y1, x2, y2], outline="red", width=6)  # Increase the width for clarity
     return image_pil
 
-# Crop a region of interest
+
 def crop_image(image, x, y, size=224):
     image_np = np.array(image)
     h, w, _ = image_np.shape
@@ -110,6 +112,7 @@ def crop_image(image, x, y, size=224):
 # Gradio UI
 with gr.Blocks() as demo:
     gr.Markdown("## Coral Classification with BeIT Model")
+    
     with gr.Row():
         with gr.Column():
             image_input = gr.Image(type="pil", label="Upload Image", interactive=True)
@@ -120,20 +123,29 @@ with gr.Blocks() as demo:
             cropped_image = gr.Image(label="Cropped Patch")
             label_output = gr.Textbox(label="Predicted Label")
     
-    # Interactions
-    def update_selection(image, x, y):
-        overlay_image = draw_rectangle(image, x, y)
-        cropped = crop_image(image, x, y)
-        return overlay_image, cropped
+    # Crop and Predict buttons
+    crop_button = gr.Button("Crop")
+    predict_button = gr.Button("Predict")
 
-    def predict_from_cropped(cropped):
-        return predict_label(cropped)
+    # Example table
+    def load_example(example_path):
+        return Image.open(example_path).convert("RGB")
 
-    crop_button = gr.Button("Crop")
-    crop_button.click(fn=update_selection, inputs=[image_input, x_slider, y_slider], outputs=[interactive_image, cropped_image])
+        # Generate table of examples
+    with gr.Row():
+        gr.Markdown("### Example Images for Quick Testing")
 
-    predict_button = gr.Button("Predict")
-    predict_button.click(fn=predict_from_cropped, inputs=cropped_image, outputs=label_output)
+    with gr.Row():
+        for genus, path in example_images.items():
+            with gr.Column():
+                thumbnail = gr.Image(value=path, interactive=False, label=genus)
+                select_button = gr.Button(value=f"Select {genus}")
+                select_button.click(fn=lambda p=path: load_example(p), inputs=None, outputs=image_input)
+
+    # Button functionality
+    crop_button.click(fn=lambda img, x, y: (draw_rectangle(img, x, y), crop_image(img, x, y)), 
+                      inputs=[image_input, x_slider, y_slider], outputs=[interactive_image, cropped_image])
+    predict_button.click(fn=predict_label, inputs=cropped_image, outputs=label_output)
 
     def update_sliders(image):
         if image:
@@ -143,4 +155,5 @@ with gr.Blocks() as demo:
 
     image_input.change(fn=update_sliders, inputs=image_input, outputs=[x_slider, y_slider])
 
+# demo.launch()
 demo.launch(server_name="0.0.0.0", server_port=7860)