Update app.py

app.py

@@ -5,10 +5,11 @@ from transformers import pipeline
 from PIL import Image
 import base64
 import io
+import requests
 
 app = FastAPI(title="STOA Chest X-Ray API")
 
-# Allow the frontend and OpenClaw to communicate with this node
+# --- CORS ---
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],
@@ -17,14 +18,18 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-# Load the ViT model into RAM (Downloads instantly because we cache it in Docker)
+# --- MODEL LOADING ---
 print("Booting Pulmonology Agent. Loading ViT model into memory...")
 pipe = pipeline("image-classification", model="dima806/chest_xray_pneumonia_detection")
 print("Agent Ready!")
 
-# Define the expected JSON payload from OpenClaw
+# --- REQUEST SCHEMA ---
+# Pydantic allows either field to be optional, so the user can send one or the other
 class PredictRequest(BaseModel):
-    image: str
+    image: str | None = None
+    image_url: str | None = None
+
+# --- ENDPOINTS ---
 
 @app.get("/health")
 def health_check():
@@ -33,19 +38,31 @@ def health_check():
 @app.post("/predict")
 def predict(req: PredictRequest):
     try:
-        # 1. Clean the base64 string if it contains the HTML data URI prefix
-        b64_data = req.image
-        if "," in b64_data:
-            b64_data = b64_data.split(",")[1]
-            
-        # 2. Convert base64 string back into raw image bytes
-        image_bytes = base64.b64decode(b64_data)
-        img = Image.open(io.BytesIO(image_bytes)).convert("RGB")
+        img = None
         
-        # 3. Execute Neural Network Math
+        # 1. Handle URL Input
+        if req.image_url:
+            response = requests.get(req.image_url, stream=True)
+            if response.status_code != 200:
+                raise Exception("Could not download image from URL.")
+            img = Image.open(response.raw).convert("RGB")
+            
+        # 2. Handle Base64 Input
+        elif req.image:
+            b64_data = req.image
+            if "," in b64_data:
+                b64_data = b64_data.split(",")[1]
+            image_bytes = base64.b64decode(b64_data)
+            img = Image.open(io.BytesIO(image_bytes)).convert("RGB")
+            
+        # 3. Handle Empty Request
+        else:
+            raise HTTPException(status_code=400, detail="Must provide 'image' (base64) or 'image_url'.")
+
+        # 4. Execute AI Math
         results = pipe(img)
         
-        # 4. Format the output to exactly match the Task 24 Spec
+        # 5. Format to exact Task 24 specifications
         top_pred = max(results, key=lambda x: x['score'])
         scores_dict = {res['label']: round(res['score'], 4) for res in results}