Update app.py

app.py

@@ -8,33 +8,28 @@ from PIL import Image
 import uvicorn
 
 # --- 1. إعداد التطبيق والموديلات ---
-app = FastAPI(title="YOLO + GIT Large Captioning API")
+app = FastAPI(title="YOLO + GIT Large (Color & Shape) API")
 
-# تحديد الجهاز
 device = "cuda" if torch.cuda.is_available() else "cpu"
-
-# مسار الموديل الخاص بك
 MY_MODEL_PATH = 'best.pt'
 
-print(f"🔄 جاري تحميل الموديلات على جهاز: {device}... يرجى الانتظار")
+print(f"🔄 جاري تحميل الموديلات على: {device}...")
 
-# تحميل موديل YOLO
 try:
     detection_model = YOLO(MY_MODEL_PATH)
-    print("✅ تم تحميل موديل YOLO الخاص بك بنجاح")
+    print("✅ تم تحميل YOLO بنجاح")
 except Exception as e:
-    print(f"⚠️ فشل تحميل {MY_MODEL_PATH}، سيتم استخدام الموديل الافتراضي: {e}")
+    print(f"⚠️ فشل تحميل الموديل الخاص، استخدام الافتراضي: {e}")
     detection_model = YOLO("yolov8n.pt")
 
-# --- التغيير هنا: استخدام microsoft/git-large ---
+# تحميل موديل 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", "model": "GIT-Large", "instruction": "Add /docs to the URL to test"}
+    return {"status": "Online", "mode": "Detailed Description (Color/Shape)"}
 
 # --- 2. وظيفة المعالجة ---
 
@@ -43,7 +38,6 @@ async def analyze_image(file: UploadFile = File(...)):
     data = await file.read()
     original_image = Image.open(io.BytesIO(data)).convert("RGB")
 
-    # 1. الكشف باستخدام YOLO
     results = detection_model(original_image, conf=0.25)
     integrated_results = []
 
@@ -53,43 +47,48 @@ async def analyze_image(file: UploadFile = File(...)):
             conf_score = float(box.conf[0])
             coords = box.xyxy[0].tolist() 
 
-            # 2. عملية القص (Cropping) 
-            # إضافة هامش بسيط (Padding) للقص يحسن أحياناً من وصف الموديل
-            cropped_img = original_image.crop((coords[0], coords[1], coords[2], coords[3]))
-
-            # 3. وصف الجزء المقصوص عبر موديل GIT Large
-            inputs = processor(images=cropped_img, return_tensors="pt").to(device)
+            # قص العنصر مع إضافة هامش بسيط (10 بكسل) لرؤية الألوان المحيطة والحواف بشكل أفضل
+            pad = 10
+            cropped_img = original_image.crop((
+                max(0, coords[0]-pad), 
+                max(0, coords[1]-pad), 
+                min(original_image.width, coords[2]+pad), 
+                min(original_image.height, coords[3]+pad)
+            ))
+
+            # --- التعديل الجوهري هنا: توجيه الموديل لوصف اللون والشكل ---
+            # نضع نصاً توجيهياً (Prompt) ليقوم الموديل بتكملته
+            prompt = f"a photo of a {label}, describing its color and shape:"
+            inputs = processor(images=cropped_img, text=prompt, return_tensors="pt").to(device)
             
-            # ضبط البارامترات للحصول على أفضل وصف من نسخة Large
             generated_ids = caption_model.generate(
-                pixel_values=inputs.pixel_values, 
-                max_length=50,
-                num_beams=4 # استخدام beam search يحسن الجودة في نسخة Large
+                pixel_values=inputs.pixel_values,
+                input_ids=inputs.input_ids, # تمرير البرومبت للموديل
+                max_length=60,
+                num_beams=5,
+                repetition_penalty=1.2 # لمنع تكرار الكلمات
             )
-            detailed_desc = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
+            
+            # فك التشفير
+            full_desc = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
+            
+            # تنظيف الوصف (إزالة البرومبت من البداية إذا ظهر)
+            clean_desc = full_desc.replace(prompt, "").strip()
 
             integrated_results.append({
                 "object_id": i + 1,
                 "label": label,
                 "confidence": f"{conf_score:.2f}",
-                "description": detailed_desc
+                "visual_description": clean_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=50)
-        general_desc = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
-        return {
-            "message": "No specific objects detected. General description provided.",
-            "general_description": general_desc
-        }
-
-    return {
-        "detected_count": len(integrated_results),
-        "results": integrated_results
-    }
-
-# --- 3. تشغيل السيرفر (تصحيح الشرطات السفلية) ---
-if name == "__main__":
+        return {"message": "No objects detected", "general_description": processor.batch_decode(generated_ids, skip_special_tokens=True)[0]}
+
+    return {"detected_count": len(integrated_results), "results": integrated_results}
+
+# --- 3. تشغيل السيرفر (تم تصحيح الـ Syntax) ---
+if __name__ == "__main__":
     uvicorn.run(app, host="0.0.0.0", port=7860)