new fastapi app

app.py

@@ -1,143 +1,242 @@
-import gradio as gr
-import pandas as pd
-import requests
-from PIL import Image
+from fastapi import FastAPI, UploadFile, File, HTTPException
+from pydantic import BaseModel
+import uvicorn
 import io
+import logging
+import datetime
+import re
 import os
-
-from huggingface_hub import snapshot_download
+import requests
+import pandas as pd
+from PIL import Image
 from autogluon.multimodal import MultiModalPredictor
+from huggingface_hub import snapshot_download
 
+###############################################################################
+# Logging configuration (optional)
+###############################################################################
+log_filename = "model_predictions.log"
+logging.basicConfig(
+    filename=log_filename, 
+    level=logging.INFO, 
+    format='%(asctime)s - %(message)s'
+)
 
-###########################
-# Load Model from HF Hub
-###########################
+###############################################################################
+# Model loading
+###############################################################################
 def load_model():
-    """Download your model files from the Hub and load with AutoGluon."""
+    """
+    Downloads the model from the specified huggingface hub repo and 
+    loads it using MultiModalPredictor.
+    """
     repo_id = "Honey-Bee-Society/honeybee_ml_v1"
     local_dir = snapshot_download(repo_id)
+
+    assets_path = os.path.join(local_dir, "assets.json")
+    model_checkpoint = os.path.join(local_dir, "model.ckpt")
+
+    if not os.path.exists(assets_path) or not os.path.exists(model_checkpoint):
+        raise FileNotFoundError("Required model files not found in the downloaded directory.")
+
     predictor = MultiModalPredictor.load(local_dir)
     return predictor
 
-predictor = load_model()  # Load once at startup
-
-
-###########################
-# Utility Functions
-###########################
-def resize_image_if_large(image, max_size_mb=1):
-    """Resizes the image if it is larger than `max_size_mb` MB."""
-    img_bytes = io.BytesIO()
-    image.save(img_bytes, format='PNG')
-    size_mb = len(img_bytes.getvalue()) / (1024 * 1024)
-    if size_mb > max_size_mb:
-        scale_factor = (max_size_mb / size_mb) ** 0.5
-        new_w = int(image.width * scale_factor)
-        new_h = int(image.height * scale_factor)
-        image = image.resize((new_w, new_h))
-    return image
-
-def classify_image(image: Image.Image):
+###############################################################################
+# Image processing and prediction routines
+###############################################################################
+def resize_image_proportionally(image, max_size_mb=1):
     """
-    Given a PIL Image, predict Honey Bee vs. Bumblebee vs. Vespidae.
-    Returns a dictionary of probabilities + predicted label.
+    If the in-memory size of the image is > max_size_mb, 
+    resize it proportionally.
     """
-    # Optionally resize large images
-    image = resize_image_if_large(image, 1)
+    img_byte_array = io.BytesIO()
+    image.save(img_byte_array, format='PNG')
+    img_size = len(img_byte_array.getvalue()) / (1024 * 1024)
 
-    # Convert to bytes for the predictor
-    buf = io.BytesIO()
-    image.save(buf, format='PNG')
-    df = pd.DataFrame({"image": [buf.getvalue()]})
+    if img_size > max_size_mb:
+        scale_factor = (max_size_mb / img_size) ** 0.5
+        new_width = int(image.width * scale_factor)
+        new_height = int(image.height * scale_factor)
+        image = image.resize((new_width, new_height))
+
+    return image
 
+
+def predict_image(image: Image.Image, predictor: MultiModalPredictor):
+    """
+    Run the prediction via the AutoGluon MultiModalPredictor. 
+    Returns probability dataframe for each class.
+    """
+    img_byte_array = io.BytesIO()
+    image.save(img_byte_array, format='PNG')
+    img_data = img_byte_array.getvalue()
+    df = pd.DataFrame({"image": [img_data]})
     probabilities = predictor.predict_proba(df, realtime=True)
+    return probabilities
 
-    # Adjust indices if your model’s labeling is different
+
+def determine_label(probabilities):
+    """
+    Given the probabilities DataFrame, compute the final label.
+    Returns a dict with numeric scores and a text label.
+    """
     honeybee_score = float(probabilities[1].iloc[0]) * 100
     bumblebee_score = float(probabilities[2].iloc[0]) * 100
     vespidae_score = float(probabilities[3].iloc[0]) * 100
 
     highest_score = max(honeybee_score, bumblebee_score, vespidae_score)
     if highest_score < 80:
-        label = "No bee detected (scores too low)."
+        prediction_label = "No bee detected (scores too low)."
     else:
         if honeybee_score == highest_score:
-            label = "Honey Bee"
+            prediction_label = "Honey Bee"
         elif bumblebee_score == highest_score:
-            label = "Bumblebee"
+            prediction_label = "Bumblebee"
         else:
-            label = "Vespidae (wasp/hornet)"
+            prediction_label = "Vespidae (wasp/hornet)"
 
     return {
         "honeybee_score": honeybee_score,
         "bumblebee_score": bumblebee_score,
         "vespidae_score": vespidae_score,
-        "prediction_label": label
+        "prediction_label": prediction_label
     }
 
 
-###########################
-# The Main Predict Function
-###########################
-def predict(uploaded_image_dict, fallback_url):
+def log_predictions(honeybee_score, bumblebee_score, vespidae_score, source_info):
     """
-    1) If `uploaded_image_dict["path"]` looks like a URL (http/https), 
-       fetch the image from the internet.
-    2) Otherwise, Gradio will have downloaded it to a local file path 
-       and `uploaded_image_dict["name"]` is that path. (or "path" is local)
-    3) If the user provides nothing in the first input, try `fallback_url`.
+    Log predictions to a file (optional).
     """
-    # Case 1: We have an "uploaded_image_dict" from the first input
-    if uploaded_image_dict and "path" in uploaded_image_dict:
-        path_val = uploaded_image_dict["path"]
-        if path_val.startswith("http"):
-            # It's a remote URL
-            try:
-                resp = requests.get(path_val.strip(), timeout=10)
-                resp.raise_for_status()
-                image = Image.open(io.BytesIO(resp.content))
-            except Exception as e:
-                return {"error": f"Failed to download/open image from URL: {e}"}
-        else:
-            # It's likely a local file path Gradio prepared
-            try:
-                image = Image.open(path_val)
-            except Exception as e:
-                return {"error": f"Failed to open local file: {e}"}
-    # Case 2: No image dict, but maybe we have a fallback URL
-    elif fallback_url.strip():
+    logging.info(
+        f"Source: {source_info}, "
+        f"Honeybee: {honeybee_score:.2f}%, "
+        f"Bumblebee: {bumblebee_score:.2f}%, "
+        f"Vespidae: {vespidae_score:.2f}%"
+    )
+
+###############################################################################
+# Request models
+###############################################################################
+class ImageUrlRequest(BaseModel):
+    image_url: str
+
+###############################################################################
+# FastAPI app and endpoints
+###############################################################################
+app = FastAPI(title="Honey Bee Classification API")
+
+# Load the model at startup (only once).
+predictor = load_model()
+
+@app.get("/ping")
+def ping():
+    """
+    A simple endpoint to check if the API is running.
+    """
+    return {"message": "pong"}
+
+@app.post("/predict")
+async def predict_endpoint(
+    image_url_req: ImageUrlRequest = None,
+    file: UploadFile = File(None)
+):
+    """
+    Accepts either a JSON body with `image_url` or a multipart form-data `file`.
+    Returns JSON with honeybee, bumblebee, vespidae scores, and a predicted label.
+    """
+    # 1) If user provided an image URL
+    if image_url_req and image_url_req.image_url:
+        image_url = image_url_req.image_url
+        # Download the image
+        try:
+            response = requests.get(
+                image_url, 
+                headers={"User-Agent": "HoneyBeeClassification/1.0 (+https://example.com)"}
+            )
+            if response.status_code != 200:
+                raise HTTPException(
+                    status_code=400, 
+                    detail=f"Failed to retrieve image from {image_url}. HTTP {response.status_code}"
+                )
+        except Exception as e:
+            raise HTTPException(
+                status_code=400, 
+                detail=f"Error downloading image from {image_url}: {e}"
+            )
+        
+        image_bytes = response.content
+        image_size_mb = len(image_bytes) / (1024*1024)
+        if image_size_mb > 10:
+            raise HTTPException(
+                status_code=413,
+                detail=f"Image size {image_size_mb:.2f}MB exceeds 10MB limit."
+            )
+        # Convert to PIL Image
         try:
-            resp = requests.get(fallback_url.strip(), timeout=10)
-            resp.raise_for_status()
-            image = Image.open(io.BytesIO(resp.content))
+            image = Image.open(io.BytesIO(image_bytes))
         except Exception as e:
-            return {"error": f"Failed to download/open fallback URL: {e}"}
+            raise HTTPException(
+                status_code=400, 
+                detail=f"Could not open image: {e}"
+            )
+        
+        # 2) If user instead provided a file
+    elif file is not None:
+        # Check file size
+        file_size = 0
+        file.file.seek(0, 2)  # move to end
+        file_size = file.file.tell()
+        file.file.seek(0)     # reset pointer
+        mb_size = file_size / (1024 * 1024)
+        if mb_size > 10:
+            raise HTTPException(
+                status_code=413,
+                detail=f"Uploaded file size {mb_size:.2f}MB exceeds 10MB limit."
+            )
+        
+        # Convert to PIL Image
+        try:
+            contents = await file.read()
+            image = Image.open(io.BytesIO(contents))
+        except Exception as e:
+            raise HTTPException(
+                status_code=400, 
+                detail=f"Could not open uploaded image: {e}"
+            )
+        source_info = f"uploaded_file:{file.filename}"
     else:
-        return {"error": "No image URL or fallback URL provided."}
-
-    # Classify the resulting image
-    return classify_image(image)
-
-
-###########################
-# Gradio Interface
-###########################
-#  1) First input = an Image (filepath)
-#  2) Second input = a Textbox
-# This order means data[0] is the dictionary for the image, 
-# data[1] is the string for fallback_url.
-
-demo = gr.Interface(
-    fn=predict,
-    inputs=[
-        gr.Image(type="filepath", label="Public Image or Uploaded File"),
-        gr.Textbox(label="Fallback URL (optional)")
-    ],
-    outputs="json",
-    api_name="predict",
-    title="Honey Bee Image Classification",
-    description="Provide an image dict with 'path' or a fallback URL. We'll classify the bee!"
-)
-
-if __name__ == "__main__":
-    demo.launch()
+        raise HTTPException(
+            status_code=400, 
+            detail="No image provided. Supply either `image_url` or `file`."
+        )
+
+    # Resize the image if needed
+    image = resize_image_proportionally(image)
+
+    # Predict
+    try:
+        probabilities = predict_image(image, predictor)
+        results = determine_label(probabilities)
+    except Exception as e:
+        raise HTTPException(
+            status_code=500,
+            detail=f"Prediction failed: {e}"
+        )
+
+    # Optionally log predictions
+    source_name = image_url_req.image_url if (image_url_req and image_url_req.image_url) else file.filename
+    log_predictions(
+        results["honeybee_score"],
+        results["bumblebee_score"],
+        results["vespidae_score"],
+        source_info=source_name
+    )
+
+    return results
+
+
+# If running locally, uncomment to start the server via `python app.py`
+# (On Hugging Face Spaces, a separate command may be used.)
+# if __name__ == "__main__":
+#     uvicorn.run(app, host="0.0.0.0", port=7860)

requirements.txt

@@ -1,6 +1,7 @@
-Pillow
+fastapi
+uvicorn
+pillow
+requests
 pandas
 autogluon.multimodal
-huggingface_hub
-requests
-gradio
+huggingface_hub