Update app.py

app.py

@@ -1,7 +1,6 @@
 import os
 import torch
 import io
-import shutil
 from fastapi import FastAPI, File, UploadFile
 from transformers import AutoProcessor, AutoModelForCausalLM
 from ultralytics import YOLO
@@ -9,31 +8,33 @@ from PIL import Image
 import uvicorn
 
 # --- 1. إعداد التطبيق والموديلات ---
-app = FastAPI(title="YOLO + GIT Captioning API")
+app = FastAPI(title="YOLO + GIT Large Captioning API")
 
-# تحديد الجهاز (استخدام CPU للمساحات المجانية لضمان الاستقرار)
+# تحديد الجهاز
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
-# مسار الموديل الذي رفعتِيه يدوياً في القائمة
+# مسار الموديل الخاص بك
 MY_MODEL_PATH = 'best.pt'
 
-print("🔄 جاري تحميل الموديلات... يرجى الانتظار")
+print(f"🔄 جاري تحميل الموديلات على جهاز: {device}... يرجى الانتظار")
 
-# تحميل موديل YOLO الخاص بكِ
+# تحميل موديل YOLO
 try:
     detection_model = YOLO(MY_MODEL_PATH)
     print("✅ تم تحميل موديل YOLO الخاص بك بنجاح")
 except Exception as e:
-    print(f"⚠️ فشل تحميل best.pt، سيتم استخدام الموديل الافتراضي: {e}")
+    print(f"⚠️ فشل تحميل {MY_MODEL_PATH}، سيتم استخدام الموديل الافتراضي: {e}")
     detection_model = YOLO("yolov8n.pt")
 
-# تحميل موديل GIT-base (أخف وأسرع للمساحة المجانية)
-processor = AutoProcessor.from_pretrained("microsoft/git-base")
-caption_model = AutoModelForCausalLM.from_pretrained("microsoft/git-base").to(device)
+# --- التغيير هنا: استخدام microsoft/git-large ---
+model_name = "microsoft/git-large"
+processor = AutoProcessor.from_pretrained(model_name)
+caption_model = AutoModelForCausalLM.from_pretrained(model_name).to(device)
+print(f"✅ تم تحميل موديل {model_name} بنجاح")
 
 @app.get("/")
 def home():
-    return {"status": "Online", "instruction": "Add /docs to the URL to test the model"}
+    return {"status": "Online", "model": "GIT-Large", "instruction": "Add /docs to the URL to test"}
 
 # --- 2. وظيفة المعالجة ---
 
@@ -43,7 +44,7 @@ async def analyze_image(file: UploadFile = File(...)):
     original_image = Image.open(io.BytesIO(data)).convert("RGB")
 
     # 1. الكشف باستخدام YOLO
-    results = detection_model(original_image, conf=0.20)
+    results = detection_model(original_image, conf=0.25)
     integrated_results = []
 
     for r in results:
@@ -53,11 +54,18 @@ async def analyze_image(file: UploadFile = File(...)):
             coords = box.xyxy[0].tolist() 
 
             # 2. عملية القص (Cropping) 
+            # إضافة هامش بسيط (Padding) للقص يحسن أحياناً من وصف الموديل
             cropped_img = original_image.crop((coords[0], coords[1], coords[2], coords[3]))
 
-            # 3. وصف الجزء المقصوص عبر موديل GIT
+            # 3. وصف الجزء المقصوص عبر موديل GIT Large
             inputs = processor(images=cropped_img, return_tensors="pt").to(device)
-            generated_ids = caption_model.generate(pixel_values=inputs.pixel_values, max_length=40)
+            
+            # ضبط البارامترات للحصول على أفضل وصف من نسخة Large
+            generated_ids = caption_model.generate(
+                pixel_values=inputs.pixel_values, 
+                max_length=50,
+                num_beams=4 # استخدام beam search يحسن الجودة في نسخة Large
+            )
             detailed_desc = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
 
             integrated_results.append({
@@ -67,9 +75,10 @@ async def analyze_image(file: UploadFile = File(...)):
                 "description": detailed_desc
             })
 
+    # إذا لم يتم اكتشاف أجسام، وصف الصورة كاملة
     if not integrated_results:
         inputs = processor(images=original_image, return_tensors="pt").to(device)
-        generated_ids = caption_model.generate(pixel_values=inputs.pixel_values, max_length=40)
+        generated_ids = caption_model.generate(pixel_values=inputs.pixel_values, max_length=50)
         general_desc = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
         return {
             "message": "No specific objects detected. General description provided.",
@@ -81,6 +90,6 @@ async def analyze_image(file: UploadFile = File(...)):
         "results": integrated_results
     }
 
-# --- 3. تشغيل السيرفر (تم تصحيح الشرطات السفلية هنا) ---
-if __name__ == "__main__":
+# --- 3. تشغيل السيرفر (تصحيح الشرطات السفلية) ---
+if name == "__main__":
     uvicorn.run(app, host="0.0.0.0", port=7860)