fix

app _bk.py

@@ -0,0 +1,111 @@
+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)
+    with torch.no_grad():
+        outputs = model(**inputs)
+        predicted_class = torch.argmax(outputs.logits, dim=1).item()
+    return model.config.id2label[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
+    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)
+    return image_pil
+
+# Function to crop the image
+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)
+    y = min(max(y, 0), h - size)
+    cropped = image_np[y:y+size, x:x+size]
+    return Image.fromarray(cropped)
+
+# Gradio Interface
+with gr.Blocks() as demo:
+    gr.Markdown("## Test Public Models for Coral Classification")
+
+    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")
+        with gr.Column():
+            interactive_image = gr.Image(label="Interactive Image with Selection")
+            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
+    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)
+
+    # 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:
+            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

@@ -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)