Update app.py

app.py

@@ -1,82 +1,92 @@
-import shutil
+import os
 import torch
-from fastapi import FastAPI, UploadFile, File
-from PIL import Image
-from io import BytesIO
-from ultralytics import YOLO
-from huggingface_hub import hf_hub_download
+import io
+import shutil
+from fastapi import FastAPI, File, UploadFile
 from transformers import AutoProcessor, AutoModelForCausalLM
+from ultralytics import YOLO
+from PIL import Image
+import uvicorn
 
-# --- 1. إعداد تطبيق FastAPI ---
-app = FastAPI(title="Object Detection & Captioning API")
-
-# --- 2. تحميل الموديلات (يتم لمرة واحدة عند بدء التشغيل) ---
-# تحميل موديل YOLO
-model_path = hf_hub_download(
-    repo_id="GradTeam/yolov26-objectDetection",
-    filename="best.pt"
-)
-yolo_model = YOLO(model_path)
+# --- 1. إعداد التطبيق والموديلات ---
+app = FastAPI(title="YOLO + GIT Captioning")
 
-# تحميل موديل الوصف (GIT)
+# تحديد الجهاز (GPU إذا توفر، وإلا CPU)
 device = "cuda" if torch.cuda.is_available() else "cpu"
-processor = AutoProcessor.from_pretrained("microsoft/git-large")
-git_model = AutoModelForCausalLM.from_pretrained("microsoft/git-large").to(device)
 
-# --- 3. الدوال المساعدة ---
+# تحميل الموديلات مرة واحدة عند تشغيل السيرفر
+# تأكدي أن ملف best.pt موجود في نفس المجلد
+MY_MODEL_PATH = 'best.pt'
 
-def get_yolo_detections(image_path):
-    results = yolo_model(image_path)
-    objects = []
-    for r in results:
-        boxes = r.boxes.xyxy.tolist()
-        classes = r.boxes.cls.tolist()
-        for box, cls in zip(boxes, classes):
-            name = yolo_model.names[int(cls)]
-            objects.append({
-                "name": name,
-                "box": [round(coord, 2) for coord in box] # تقريب الإحداثيات
-            })
-    return objects
-
-def get_image_caption(image_path, objects):
-    image = Image.open(image_path).convert("RGB")
-    names = [obj["name"] for obj in objects]
-    
-    # بناء الـ Prompt بناءً على الأجسام المكتشفة
-    text_prompt = "Objects detected: " + ", ".join(names) if names else "Describe this image."
-    
-    inputs = processor(images=image, text=text_prompt, return_tensors="pt").to(device)
-    generated_ids = git_model.generate(**inputs, max_length=50)
-    caption = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
-    return caption
-
-# --- 4. نقطة النهاية (API Endpoint) ---
+print("🔄 جاري تحميل الموديلات... قد يستغرق ذلك دقائق")
+
+# تحميل موديل YOLO (الخاص بكِ)
+if os.path.exists(MY_MODEL_PATH):
+    detection_model = YOLO(MY_MODEL_PATH)
+    print("✅ تم تحميل موديل YOLO بنجاح")
+else:
+    # في حال لم يجد ملفك، سيستخدم موديل أساسي مؤقتاً لكي لا يتعطل السيرفر
+    detection_model = YOLO("yolov8n.pt")
+    print("⚠️ تحذير: لم يتم العثور على best.pt، تم استخدام الموديل الافتراضي.")
+
+# تحميل موديل GIT-base (أخف وأسرع للمساحات المجانية)
+processor = AutoProcessor.from_pretrained("microsoft/git-base")
+caption_model = AutoModelForCausalLM.from_pretrained("microsoft/git-base").to(device)
+
+@app.get("/")
+def home():
+    return {"status": "Running", "message": "Go to /docs to test the API"}
+
+# --- 2. وظيفة المعالجة (نفس منطق كود كولاب الخاص بكِ) ---
 
 @app.post("/analyze")
 async def analyze_image(file: UploadFile = File(...)):
-    # حفظ الملف المرفوع مؤقتاً
-    temp_path = "temp_image.jpg"
-    with open(temp_path, "wb") as buffer:
-        shutil.copyfileobj(file.file, buffer)
-    
-    try:
-        # 1. تنفيذ كشف الأجسام
-        detected_objects = get_yolo_detections(temp_path)
-        
-        # 2. تنفيذ وصف الصورة بناءً على الأجسام
-        description = get_image_caption(temp_path, detected_objects)
-        
+    data = await file.read()
+    original_image = Image.open(io.BytesIO(data)).convert("RGB")
+
+    # استخدام نفس الإعدادات التي نجحت معكِ في كولاب
+    results = detection_model(original_image, conf=0.20)
+
+    integrated_results = []
+
+    for r in results:
+        for i, box in enumerate(r.boxes):
+            label = r.names[int(box.cls)]
+            conf_score = float(box.conf[0])
+            coords = box.xyxy[0].tolist() 
+
+            # عملية القص (Cropping)
+            cropped_img = original_image.crop((coords[0], coords[1], coords[2], coords[3]))
+
+            # وصف الجزء المقصوص
+            inputs = processor(images=cropped_img, return_tensors="pt").to(device)
+            generated_ids = caption_model.generate(pixel_values=inputs.pixel_values, max_length=40)
+            detailed_desc = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
+
+            integrated_results.append({
+                "object_id": i + 1,
+                "yolo_label": label,
+                "confidence": f"{conf_score:.2f}",
+                "detailed_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)
+        general_desc = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
         return {
-            "status": "success",
-            "detected_objects_count": len(detected_objects),
-            "objects": detected_objects,
-            "description": description
+            "message": "YOLO لم يكتشف أجساماً، إليك وصفاً عاماً.",
+            "general_description": general_desc
         }
-    except Exception as e:
-        return {"status": "error", "message": str(e)}
 
-# --- 5. التشغيل (اختياري محلياً) ---
+    return {
+        "detected_count": len(integrated_results),
+        "details": integrated_results
+    }
+
+# --- 3. تشغيل السيرفر (إعدادات Hugging Face) ---
 if name == "__main__":
-    import uvicorn
+    # المنفذ 7860 هو الافتراضي والمطلوب في Hugging Face
     uvicorn.run(app, host="0.0.0.0", port=7860)
+