Update max.py

max.py

@@ -5,8 +5,10 @@ import re
 import logging
 import tempfile
 import base64
+import requests
 from uuid import uuid4
 from typing import Optional, List
+from urllib.parse import urlparse
 from fastapi import FastAPI, UploadFile, File, HTTPException
 from fastapi.responses import JSONResponse
 from fastapi.middleware.cors import CORSMiddleware
@@ -245,6 +247,9 @@ class PredictionResponse(BaseModel):
     category: str
     output_image: Optional[str] = None  # Base64-encoded output image
 
+class ImageUrlRequest(BaseModel):
+    image_url: str
+
 class QuestionRequest(BaseModel):
     session_id: str
     question: str
@@ -293,6 +298,69 @@ food_categories = {
     "Sauce Condiments and Seasonings": ['guacamole' ],
 }
 
+# Helper functions for URL processing
+def convert_google_drive_url(url: str) -> str:
+    """Convert Google Drive share URL to direct download URL."""
+    if 'drive.google.com' in url:
+        if '/file/d/' in url:
+            # Extract file ID from share URL
+            match = re.search(r'/file/d/([a-zA-Z0-9-_]+)', url)
+            if match:
+                file_id = match.group(1)
+                return f"https://drive.google.com/uc?id={file_id}&export=download"
+        elif '/uc?id=' in url:
+            # Already a direct download URL, ensure it has export=download
+            if 'export=download' not in url:
+                if '&' in url:
+                    return url + "&export=download"
+                else:
+                    return url + "?export=download"
+            return url
+        elif 'drive.usercontent.google.com' in url:
+            # Handle usercontent URLs
+            return url.replace('export=view', 'export=download')
+    
+    # Return original URL if not a Google Drive URL
+    return url
+
+def download_image_from_url(url: str) -> Image.Image:
+    """Download image from URL and return PIL Image object."""
+    try:
+        # Set headers to mimic a browser request
+        headers = {
+            'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36',
+            'Accept': 'image/webp,image/apng,image/*,*/*;q=0.8',
+            'Accept-Language': 'en-US,en;q=0.9',
+            'Accept-Encoding': 'gzip, deflate, br',
+            'Connection': 'keep-alive',
+            'Upgrade-Insecure-Requests': '1',
+        }
+        
+        # Download the image with a timeout
+        response = requests.get(url, headers=headers, timeout=30, stream=True)
+        response.raise_for_status()
+        
+        # Check content type
+        content_type = response.headers.get('content-type', '').lower()
+        if not content_type.startswith('image/'):
+            # Sometimes Google Drive doesn't return proper content-type, so check content
+            if len(response.content) < 1000 and b'html' in response.content.lower():
+                raise ValueError(f"URL returned HTML content, not an image. Please ensure the Google Drive file is publicly accessible.")
+        
+        # Open image from response content
+        image = Image.open(io.BytesIO(response.content)).convert("RGB")
+        return image
+        
+    except requests.exceptions.RequestException as e:
+        logger.error(f"Failed to download image from URL {url}: {str(e)}")
+        raise HTTPException(status_code=400, detail=f"Failed to download image: {str(e)}. Please ensure the Google Drive link is publicly accessible.")
+    except UnidentifiedImageError:
+        logger.error(f"Invalid image file from URL {url}")
+        raise HTTPException(status_code=400, detail="URL does not contain a valid image file")
+    except Exception as e:
+        logger.error(f"Error processing image from URL {url}: {str(e)}")
+        raise HTTPException(status_code=500, detail=f"Error processing image: {str(e)}")
+
 # Image classification endpoints
 @app.post(
     "/predict",
@@ -355,6 +423,70 @@ async def predict_image(file: UploadFile = File(..., description="A food image i
     logger.info(f"Prediction for {file.filename}: {label} (Category: {category})")
     return PredictionResponse(label=label, category=category, output_image=output_image)
 
+@app.post(
+    "/predict-from-url",
+    response_model=PredictionResponse,
+    summary="Classify a food image from URL",
+    description="Provide an image URL (preferably Google Drive) to classify it into one of 101 food categories and return its category."
+)
+async def predict_image_from_url(request: ImageUrlRequest):
+    image_url = request.image_url.strip()
+    logger.info(f"Received image URL for prediction: {image_url}")
+    
+    if not image_url:
+        raise HTTPException(status_code=400, detail="Image URL is required")
+    
+    # Validate URL format
+    try:
+        parsed_url = urlparse(image_url)
+        if not parsed_url.scheme or not parsed_url.netloc:
+            raise HTTPException(status_code=400, detail="Invalid URL format")
+    except Exception:
+        raise HTTPException(status_code=400, detail="Invalid URL format")
+    
+    # Convert Google Drive URL to direct download URL if needed
+    direct_url = convert_google_drive_url(image_url)
+    logger.info(f"Converted URL: {direct_url}")
+    
+    # Download and process image
+    try:
+        image = download_image_from_url(direct_url)
+    except HTTPException:
+        raise
+    except Exception as e:
+        logger.error(f"Unexpected error downloading image: {str(e)}")
+        raise HTTPException(status_code=500, detail="Failed to process image from URL")
+
+    # Predict using the downloaded image
+    try:
+        # Create a temporary file for prediction
+        fd, temp_file_path = tempfile.mkstemp(suffix=".jpg")
+        try:
+            image.save(temp_file_path)
+            label, output_image_path = classifier.predict(temp_file_path)
+        finally:
+            os.close(fd)
+            if os.path.exists(temp_file_path):
+                os.remove(temp_file_path)  # Clean up
+    except Exception as e:
+        logger.error(f"Prediction error: {str(e)}")
+        raise HTTPException(status_code=500, detail=f"Prediction error: {str(e)}")
+
+    # Determine category
+    category = next((cat for cat, foods in food_categories.items() if label in foods), "Uncategorized")
+
+    # Encode output image as base64 if available
+    output_image = None
+    if output_image_path and os.path.exists(output_image_path):
+        try:
+            with open(output_image_path, "rb") as f:
+                output_image = base64.b64encode(f.read()).decode("utf-8")
+        except Exception as e:
+            logger.warning(f"Failed to encode output image: {str(e)}")
+
+    logger.info(f"Prediction for URL {image_url}: {label} (Category: {category})")
+    return PredictionResponse(label=label, category=category, output_image=output_image)
+
 # E-commerce assistant endpoints
 def execute_function_call(raw_response: str) -> str:
     think_end = raw_response.find('</think>')
@@ -492,6 +624,7 @@ async def root():
         "message": "Welcome to the EcoHarvest Combined API",
         "endpoints": {
             "image_classification": "/predict",
+            "image_classification_from_url": "/predict-from-url",
             "ecommerce_assistant": "/ask",
             "food_categories": "/food/categories",
             "docs": "/docs"