Update app.py

app.py

@@ -1,9 +1,11 @@
-from fastapi import FastAPI, File, UploadFile, HTTPException
+from fastapi import FastAPI, File, UploadFile, HTTPException, Query
 from fastapi.responses import JSONResponse
 from tensorflow.keras.models import load_model
 from PIL import Image
 import numpy as np
 import requests
+import io
+from typing import Optional
 
 # Initialize FastAPI app
 app = FastAPI()
@@ -19,7 +21,7 @@ breed_names = {
 }
 
 # Function to preprocess the image
-def preprocess_image(image):
+def preprocess_image(image: Image.Image):
     image = image.resize((150, 150))
     img_array = np.array(image)
     img_array = img_array / 255.0
@@ -27,14 +29,14 @@ def preprocess_image(image):
     return img_array
 
 # Function to classify the breed
-def classify_breed(image, model):
+def classify_breed(image: Image.Image, model):
     img_array = preprocess_image(image)
     predictions = model.predict(img_array)
     predicted_class_index = np.argmax(predictions)
-    return breed_names[predicted_class_index]
+    return breed_names.get(predicted_class_index, "Unknown")
 
 # Function to fetch breed information from an API
-def fetch_breed_info(breed_name):
+def fetch_breed_info(breed_name: str):
     url = f'https://api.thedogapi.com/v1/breeds/search?q={breed_name}'
     response = requests.get(url)
     if response.status_code == 200:
@@ -45,14 +47,25 @@ def fetch_breed_info(breed_name):
 
 # API route for prediction
 @app.post("/predict")
-async def predict(file: UploadFile = File(...)):
+async def predict(
+    file: Optional[UploadFile] = File(None),
+    url: Optional[str] = Query(None, description="Public URL of the image")
+):
     try:
-        # Check file type
-        if file.content_type not in ["image/jpeg", "image/png", "image/jpg"]:
-            raise HTTPException(status_code=400, detail="Invalid file type. Only JPG and PNG are allowed.")
-
-        # Read and process the image
-        image = Image.open(file.file)
+        # Determine input method: file has priority over URL.
+        if file is not None:
+            # Check file type
+            if file.content_type not in ["image/jpeg", "image/png", "image/jpg"]:
+                raise HTTPException(status_code=400, detail="Invalid file type. Only JPG and PNG are allowed.")
+            image = Image.open(file.file)
+        elif url is not None:
+            # Download image from URL
+            resp = requests.get(url)
+            if resp.status_code != 200:
+                raise HTTPException(status_code=400, detail="Unable to fetch image from provided URL.")
+            image = Image.open(io.BytesIO(resp.content))
+        else:
+            raise HTTPException(status_code=400, detail="No image provided. Please upload a file or provide a URL.")
 
         # Classify the breed
         breed_name = classify_breed(image, model)
@@ -63,7 +76,7 @@ async def predict(file: UploadFile = File(...)):
         # Prepare response
         response = {
             "predicted_breed": breed_name,
-            "breed_info": breed_info[0] if breed_info else "No additional information available."
+            "breed_info": breed_info[0] if breed_info and len(breed_info) > 0 else "No additional information available."
         }
         return JSONResponse(content=response)
 
@@ -73,4 +86,4 @@ async def predict(file: UploadFile = File(...)):
 # Root route
 @app.get("/")
 def read_root():
-    return {"message": "Welcome to the Dog Breed Classification API!"}
+    return {"message": "Welcome to the Dog Breed Classification API!"}