Update app.py

app.py

@@ -3,6 +3,7 @@ from pydantic import BaseModel
 from fastapi.middleware.cors import CORSMiddleware
 from transformers import pipeline
 from PIL import Image
+from bs4 import BeautifulSoup
 import base64
 import io
 import requests
@@ -24,13 +25,49 @@ pipe = pipeline("image-classification", model="dima806/chest_xray_pneumonia_dete
 print("Agent Ready!")
 
 # --- REQUEST SCHEMA ---
-# Pydantic allows either field to be optional, so the user can send one or the other
 class PredictRequest(BaseModel):
     image: str | None = None
     image_url: str | None = None
 
-# --- ENDPOINTS ---
+# --- SMART FETCHER HELPER ---
+def get_image_from_any_url(url: str):
+    """Smart fetcher that handles both raw images and HTML webpages."""
+    headers = {
+        "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36",
+        "Accept": "text/html,application/xhtml+xml,application/xml;q=0.9,image/avif,image/webp,*/*;q=0.8",
+        "Referer": "https://google.com"
+    }
+    
+    # 1. Fetch whatever is at the URL
+    response = requests.get(url, headers=headers, timeout=10)
+    if response.status_code != 200:
+        raise Exception(f"Site blocked us (Error {response.status_code})")
+        
+    content_type = response.headers.get('Content-Type', '').lower()
+    
+    # 2. If it's already an image, just return it
+    if content_type.startswith('image/'):
+        return Image.open(io.BytesIO(response.content)).convert("RGB")
+        
+    # 3. If it's a webpage, hunt for the main Open Graph image
+    elif content_type.startswith('text/html'):
+        print("Webpage detected! Scraping for the main image...")
+        soup = BeautifulSoup(response.text, 'html.parser')
+        og_image = soup.find('meta', property='og:image')
+        
+        if og_image and og_image.get('content'):
+            actual_image_url = og_image['content']
+            print(f"Found hidden image at: {actual_image_url}")
+            
+            img_response = requests.get(actual_image_url, headers=headers, timeout=10)
+            return Image.open(io.BytesIO(img_response.content)).convert("RGB")
+        else:
+            raise Exception("Could not find a main image on this webpage.")
+            
+    else:
+        raise Exception(f"Unsupported link type: {content_type}")
 
+# --- ENDPOINTS ---
 @app.get("/health")
 def health_check():
     return {"status": "ok"}
@@ -40,16 +77,9 @@ def predict(req: PredictRequest):
     try:
         img = None
         
-        # 1. Handle URL Input
+        # 1. Handle URL Input (Using the new Smart Fetcher)
         if req.image_url:
-            # The fake ID card so Wikipedia lets our AI download the image
-            headers = {"User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64)"}
-            response = requests.get(req.image_url, stream=True, headers=headers)
-            
-            if response.status_code != 200:
-                raise Exception(f"Could not download image. Server returned: {response.status_code}")
-                
-            img = Image.open(response.raw).convert("RGB")
+            img = get_image_from_any_url(req.image_url)
             
         # 2. Handle Base64 Input
         elif req.image: