Upload 3 files

app.py

@@ -1,13 +1,14 @@
 import os
+import io
 import json
 import torch
-from fastapi import FastAPI, HTTPException
+from fastapi import FastAPI, HTTPException, UploadFile, File, Form
 from pydantic import BaseModel
 from transformers import AutoProcessor, AutoModelForImageTextToText
 from fastapi.middleware.cors import CORSMiddleware
+from PIL import Image
 
 # Define the model ID
-# MedGemma 1.5 4B fits in ~8GB RAM using bfloat16, perfect for HF CPU Spaces
 MODEL_ID = "google/medgemma-1.5-4b-it"
 
 # Get huggingface token for gated models
@@ -15,7 +16,7 @@ HF_TOKEN = os.environ.get("HF_TOKEN")
 
 app = FastAPI(
     title="MedGemma Radiology API",
-    description="FastAPI service for analyzing radiology reports using MedGemma.",
+    description="FastAPI service for analyzing multimodal radiology cases (Image + Text) using MedGemma.",
     version="1.0.0"
 )
 
@@ -35,13 +36,12 @@ def load_model():
     global processor, model
     print(f"Loading processor and model {MODEL_ID}...")
     try:
-        # Check deployment environment device
         device = "cuda" if torch.cuda.is_available() else "cpu"
         
         processor = AutoProcessor.from_pretrained(MODEL_ID, token=HF_TOKEN)
         model = AutoModelForImageTextToText.from_pretrained(
             MODEL_ID,
-            torch_dtype=torch.bfloat16, # Optimized for reasonable RAM usage
+            torch_dtype=torch.bfloat16, 
             device_map=device,
             low_cpu_mem_usage=True,
             token=HF_TOKEN
@@ -52,46 +52,70 @@ def load_model():
         print(f"Error loading model: {e}")
         print("Make sure you have set the HF_TOKEN environment variable correctly and accepted the model license.")
 
-class RadiologyCase(BaseModel):
-    case_description: str
-
 class AnalysisResult(BaseModel):
     diagnosis: str
     recommendations: str
     urgency_level: str
-    raw_response: str = None # Included internally for debugging
+    raw_response: str = None
 
 # The "dماغ" or System Prompt
-SYSTEM_PROMPT = """أنت الآن "مساعد تشخيص إشعاعي ذكي" متطور. مهمتك هي تحليل النصوص الواردة إليك والتي تصف نتائج صور الأشعة (X-ray, CT, MRI).
+SYSTEM_PROMPT = """أنت الآن "مساعد تشخيص إشعاعي ذكي" متطور. مهمتك هي تحليل الصور والفحوصات الطبية المرفقة بالإضافة إلى النصوص الواردة والتي تصف حالة المريض.
 
 قواعد العمل:
-1. التخصص: ركز فقط على المصطلحات الطبية الإشعاعية (مثل Opacity, Radiolucency, Fracture, Lesion).
-2. الهيكلية: يجب أن يكون ردك منظماً (النتائج الأساسية، التشخيص المحتمل، التوصيات).
-3. الدقة: إذا كانت الحالة طارئة (مثل نزيف أو كسر مضاعف)، ابدأ بردك واجعل مستوى الحالة "حالة طارئة - Urgent".
+1. التخصص: ركز فقط على المصطلحات الطبية الإشعاعية (مثل Opacity, Radiolucency, Fracture, Lesion) عند وصف الصورة.
+2. الهيكلية: يجب أن يكون ردك منظماً (النتائج الأساسية للصورة، التشخيص المحتمل، التوصيات).
+3. الدقة: إذا كانت الحالة طارئة بناءً على الصورة (مثل كسر مضاعف أو استرواح الصدر)، اجعل مستوى الحالة "حالة طارئة - Urgent".
 4. التحذير: أضف دائماً في التوصيات أن هذا التحليل هو "رأي استشاري ذكي" ويجب مراجعته من قبل طبيب أشعة مختص.
 5. اللغة: أجب باللغة العربية الطبية الرصينة.
 
 مهم جداً: قم بالرد باستخدام صيغة JSON صحيحة تحتوي على المفاتيح التالية فقط:
 {
-  "diagnosis": "التشخيص المحتمل والنتائج الأساسية",
+  "diagnosis": "نتائج تحليل الصورة والتشخيص المحتمل",
   "recommendations": "التوصيات والتحذير",
   "urgency_level": "مستوى الحالة (مثلاً: حالة طارئة - Urgent أو عادية - Normal)"
 }"""
 
 @app.post("/analyze-radiology", response_model=AnalysisResult)
-async def analyze_report(case: RadiologyCase):
+async def analyze_report(
+    case_description: str = Form(""), 
+    image: UploadFile = File(None)
+):
+    """
+    Analyzes a radiology case. Accepts an optional text description and an optional image (X-Ray, MRI, etc).
+    At least one of them must be provided.
+    """
     if not model or not processor:
         raise HTTPException(status_code=503, detail="The AI model is currently loading or failed to load. Please try again later.")
         
+    if not case_description and not image:
+        raise HTTPException(status_code=400, detail="يجب إرفاق صورة أو كتابة وصف للحالة على الأقل.")
+
     try:
-        # Combine System prompt with user case
-        user_text = f"{SYSTEM_PROMPT}\n\nنص التقرير أو الحالة:\n{case.case_description}"
+        content = []
+        
+        # 1. Process Image if provided
+        if image:
+            image_data = await image.read()
+            pil_image = Image.open(io.BytesIO(image_data)).convert("RGB")
+            content.append({"type": "image", "image": pil_image})
+            
+        # 2. Process Text
+        user_text = SYSTEM_PROMPT + "\n\n"
+        if case_description:
+            user_text += f"وصف الحالة السريرية أو الأعراض:\n{case_description}\n\n"
+        
+        if image:
+            user_text += "الرجاء تحليل الصورة الطبية المرفقة بناءً على القواعد أعلاه."
+        else:
+            user_text += "الرجاء تحليل الوصف الطبي أعلاه بناءً على القواعد أعلاه."
+            
+        content.append({"type": "text", "text": user_text})
+
+        # 3. Create messages format
         messages = [
             {
                 "role": "user",
-                "content": [
-                    {"type": "text", "text": user_text}
-                ]
+                "content": content
             }
         ]
         
@@ -103,22 +127,21 @@ async def analyze_report(case: RadiologyCase):
 
         input_len = inputs["input_ids"].shape[-1]
         
-        # Generate with optimized settings
+        # Generate
         with torch.inference_mode():
             generation = model.generate(
                 **inputs, 
                 max_new_tokens=1024, 
                 do_sample=True, 
-                temperature=0.2, # Conservative temp for medical accuracy
+                temperature=0.2,
                 top_p=0.9
             )
-            # Exclude the input prompt from generation output
             generation_output = generation[0][input_len:]
             
         decoded = processor.decode(generation_output, skip_special_tokens=True)
         raw_output = decoded.strip()
         
-        # Helper: Clean out markdown block delimiters if model generated them
+        # Clean JSON markdown blocks
         clean_json = raw_output
         if clean_json.startswith("```json"):
             clean_json = clean_json.replace("```json", "", 1)
@@ -130,7 +153,6 @@ async def analyze_report(case: RadiologyCase):
         try:
             parsed_data = json.loads(clean_json)
         except json.JSONDecodeError:
-            # Fallback if model doesn't strictly adhere to JSON outline
             is_urgent = "Urgent" in raw_output or "طارئة" in raw_output
             parsed_data = {
                 "diagnosis": raw_output[:500] + ("..." if len(raw_output)>500 else ""),
@@ -153,5 +175,6 @@ def health_check():
     return {
         "status": "Online",
         "model": MODEL_ID,
-        "message": "Welcome to MedGemma Radiology API"
+        "vision_enabled": True,
+        "message": "Welcome to Multimodal MedGemma Radiology API"
     }

requirements.txt

@@ -1,6 +1,7 @@
 fastapi>=0.104.1
 uvicorn>=0.23.2
 pydantic>=2.4.2
+python-multipart>=0.0.9
 torch>=2.1.0
 transformers>=4.40.0
 accelerate>=0.29.3