Update app.py

app.py

@@ -13,15 +13,19 @@ sys.modules["flash_attn.bert_padding"] = types.ModuleType("flash_attn.bert_paddi
 
 import io
 import time
+import httpx
 import torch
 from PIL import Image
 from transformers import AutoProcessor, AutoModelForCausalLM
 from fastapi import FastAPI, HTTPException, UploadFile, File
+from fastapi.middleware.cors import CORSMiddleware
+from fastapi.responses import JSONResponse
+from pydantic import BaseModel
 from contextlib import asynccontextmanager
+from typing import Optional
 
 MODEL_ID = "microsoft/Florence-2-large-ft"
 
-# ─── السؤال الأصلي + تأكيد على اليد ─────────────────────────────
 VQA_QUESTION = (
     "Is there a woman or any part of a woman's body in this image? Answer yes or no only."
 )
@@ -46,82 +50,107 @@ async def lifespan(app: FastAPI):
     MODEL_DATA.clear()
 
 app = FastAPI(
-    title="Female Detection API - VQA",
-    description="Florence-2-large-ft | VQA",
-    version="4.3.0",
+    title="AI Shield - Female Detection API",
+    description="Florence-2-large-ft | VQA | Compatible with AI Shield Chrome Extension",
+    version="5.0.0",
     lifespan=lifespan
 )
 
-@app.get("/health")
-def health():
-    return {"status": "ok", "model_loaded": "model" in MODEL_DATA}
+# ─── CORS: ضروري لإضافة Chrome ───────────────────────────────────
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+# ─── Schemas ─────────────────────────────────────────────────────
+class ImageUrlRequest(BaseModel):
+    image_url: str          # من إضافة Chrome
+
+# ─── دالة التحليل المشتركة ────────────────────────────────────────
+def analyze_image(image: Image.Image) -> dict:
+    processor = MODEL_DATA["processor"]
+    model     = MODEL_DATA["model"]
+
+    task   = "<VQA>"
+    prompt = f"{task}{VQA_QUESTION}"
+
+    inputs = processor(text=prompt, images=image, return_tensors="pt")
 
-def decide(answer: str) -> tuple[str, str]:
-    """
-    - "no"        → allow  ✅
-    - "yes"       → block  🔴
-    - أي شيء آخر → block  🔴 للأمان
-    """
-    a = answer.strip().lower()
+    start_time = time.time()
+    with torch.no_grad():
+        generated_ids = model.generate(
+            input_ids=inputs["input_ids"],
+            pixel_values=inputs["pixel_values"],
+            max_new_tokens=10,
+            num_beams=3,
+            do_sample=False
+        )
+
+    generated_text = processor.batch_decode(generated_ids, skip_special_tokens=False)[0]
+    parsed  = processor.post_process_generation(
+        generated_text,
+        task=task,
+        image_size=(image.width, image.height)
+    )
+    elapsed = round(time.time() - start_time, 2)
+    answer  = parsed.get(task, "").strip()
+
+    # ─── منطق القرار ─────────────────────────────────────────────
+    a = answer.lower()
     if a == "no" or a.startswith("no"):
-        return "allow", "model_answered_no"
+        decision, reason = "ALLOW", "model_answered_no"
     elif "yes" in a:
-        return "block", "model_answered_yes"
+        decision, reason = "BLOCK", "model_answered_yes"
     else:
-        return "block", "unexpected_answer_blocked_for_safety"
+        decision, reason = "BLOCK", "unexpected_answer_blocked_for_safety"
 
-@app.post("/analyze")
-async def analyze_image(file: UploadFile = File(...)):
+    return {
+        "decision":       decision,   # ALLOW | BLOCK (بالحروف الكبيرة لتتوافق مع الإضافة)
+        "reason":         reason,
+        "vqa_answer":     answer,
+        "execution_time": elapsed,
+        "status":         "success"
+    }
 
-    if not file.content_type.startswith("image/"):
-        raise HTTPException(status_code=400, detail="الملف ليس صورة")
+# ─── Health Check ─────────────────────────────────────────────────
+@app.get("/health")
+def health():
+    return {"status": "ok", "model_loaded": "model" in MODEL_DATA}
 
+# ─── Endpoint 1: من إضافة Chrome (image_url) ─────────────────────
+# background.js يرسل: POST /analyze {"image_url": "https://..."}
+@app.post("/analyze")
+async def analyze_from_url(request: ImageUrlRequest):
     try:
-        image = Image.open(io.BytesIO(await file.read())).convert("RGB")
+        async with httpx.AsyncClient(timeout=30) as client:
+            response = await client.get(request.image_url)
+            response.raise_for_status()
+            image_bytes = response.content
     except Exception as e:
-        raise HTTPException(status_code=400, detail=f"خطأ في قراءة الصورة: {str(e)}")
+        raise HTTPException(status_code=400, detail=f"فشل تحميل الصورة من URL: {str(e)}")
 
     try:
-        processor = MODEL_DATA["processor"]
-        model     = MODEL_DATA["model"]
-
-        task   = "<VQA>"
-        prompt = f"{task}{VQA_QUESTION}"
-
-        inputs = processor(text=prompt, images=image, return_tensors="pt")
-
-        start_time = time.time()
-        with torch.no_grad():
-            generated_ids = model.generate(
-                input_ids=inputs["input_ids"],
-                pixel_values=inputs["pixel_values"],
-                max_new_tokens=10,
-                num_beams=3,
-                do_sample=False
-            )
-
-        generated_text = processor.batch_decode(generated_ids, skip_special_tokens=False)[0]
-        parsed  = processor.post_process_generation(
-            generated_text,
-            task=task,
-            image_size=(image.width, image.height)
-        )
-        elapsed = round(time.time() - start_time, 2)
+        image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
+    except Exception as e:
+        raise HTTPException(status_code=400, detail=f"خطأ في قراءة الصورة: {str(e)}")
 
-        answer           = parsed.get(task, "").strip()
-        decision, reason = decide(answer)
+    return analyze_image(image)
 
-        return {
-            "decision":       decision,
-            "reason":         reason,
-            "vqa_answer":     answer,
-            "question":       VQA_QUESTION,
-            "execution_time": elapsed,
-            "status":         "success"
-        }
+# ─── Endpoint 2: اختبار يدوي (رفع ملف) ───────────────────────────
+@app.post("/analyze-file")
+async def analyze_from_file(file: UploadFile = File(...)):
+    if not file.content_type.startswith("image/"):
+        raise HTTPException(status_code=400, detail="الملف ليس صورة")
 
+    try:
+        image = Image.open(io.BytesIO(await file.read())).convert("RGB")
     except Exception as e:
-        raise HTTPException(status_code=500, detail=str(e))
+        raise HTTPException(status_code=400, detail=f"خطأ في قراءة الصورة: {str(e)}")
+
+    return analyze_image(image)
 
 
 if __name__ == "__main__":