fix

app.py

@@ -12,18 +12,24 @@ MODEL_LIST = {
     'convnext': "facebook/convnext-tiny-224",
 }
 
-# Preprocessing transforms
-def get_preprocessor(model_name):
-    extractor = AutoFeatureExtractor.from_pretrained(MODEL_LIST[model_name])
-    return extractor
+# Global variables
+current_model = None
+current_preprocessor = None
 
-# Load a model from Hugging Face
-def load_model(model_name):
-    model = AutoModelForImageClassification.from_pretrained(MODEL_LIST[model_name]).cuda().eval()
-    return model
+# 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}")
+    current_model = AutoModelForImageClassification.from_pretrained(MODEL_LIST[model_name]).cuda().eval()
+    current_preprocessor = AutoFeatureExtractor.from_pretrained(MODEL_LIST[model_name])
+    return f"Model {model_name} loaded successfully."
 
-# Function to make predictions
+# 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("cuda")
     with torch.no_grad():
         outputs = model(**inputs)
@@ -32,6 +38,7 @@ def predict(image, model, preprocessor):
 
 # 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
@@ -41,6 +48,7 @@ def draw_rectangle(image, x, y, size=224):
 
 # 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)
@@ -48,10 +56,6 @@ def crop_image(image, x, y, size=224):
     cropped = image_np[y:y+size, x:x+size]
     return Image.fromarray(cropped)
 
-# Global variables
-current_model = None
-current_preprocessor = None
-
 # Gradio Interface
 with gr.Blocks() as demo:
     gr.Markdown("## Test Public Models for Coral Classification")
@@ -67,12 +71,9 @@ with gr.Blocks() as demo:
             cropped_image = gr.Image(label="Cropped Patch")
             label_output = gr.Textbox(label="Predicted Label")
 
-    # Update the current model and preprocessor
+    # Update the model and preprocessor
     def update_model(model_name):
-        global current_model, current_preprocessor
-        current_model = load_model(model_name)
-        current_preprocessor = get_preprocessor(model_name)
-        return f"Model {model_name} loaded successfully."
+        return load_model_and_preprocessor(model_name)
 
     # Update the rectangle and crop the patch
     def update_selection(image, x, y):
@@ -82,6 +83,7 @@ with gr.Blocks() as demo:
 
     # 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
@@ -102,4 +104,7 @@ with gr.Blocks() as demo:
 
     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)