frist commit

router/image_clf.py

@@ -1,7 +1,10 @@
 import os
 import time
-from typing import Union
-from fastapi import APIRouter
+import requests
+from PIL import Image
+from io import BytesIO
+
+from fastapi import APIRouter, HTTPException
 from scripts.data_model import ImageInput, ImageOutput
 from utils.pipeline import load_model
 
@@ -10,6 +13,16 @@ router = APIRouter()
 BASE_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))  
 MODEL_PATH = os.path.join(BASE_DIR, "ml-models", "vit-human-pose-classification/")
 
+def download_image(url):
+    """Download and open an image from a URL."""
+    try:
+        response = requests.get(url)
+        response.raise_for_status()
+        return Image.open(BytesIO(response.content)).convert("RGB")
+    except Exception as e:
+        raise HTTPException(status_code=400, detail=f"Failed to download image: {e}")
+
+
 @router.post(
     "/image_classification",
     response_model=ImageOutput,
@@ -19,11 +32,13 @@ MODEL_PATH = os.path.join(BASE_DIR, "ml-models", "vit-human-pose-classification/
 def image_classification(input: ImageInput)-> ImageOutput:
     try:
         pipe = load_model(MODEL_PATH, is_image_model=True)
+        image = download_image(input.url)
+        
         start = time.time()
-        output = pipe(input.url)
+        output = pipe(image)
         end = time.time()
+        
         prediction_time = int((end-start)*1000)
-
         labels_and_scores = [{"label": x['label'], "score": x['score']} for x in output]
 
         return ImageOutput(