app.py

# app.py

import os
import json
import traceback
import torch
import gradio as gr
import numpy as np
from PIL import Image
import cv2
import math

# --- استيراد من الملفات المنظمة في مشروعك ---
from model import build_interfuser_model 
from logic import (
    transform, lidar_transform, InterfuserController, ControllerConfig,
    Tracker, DisplayInterface, render, render_waypoints, render_self_car,
    ensure_rgb, WAYPOINT_SCALE_FACTOR, T1_FUTURE_TIME, T2_FUTURE_TIME
)

# ==============================================================================
#           1. إعدادات ومسارات النماذج
# ==============================================================================
WEIGHTS_DIR = "model"
EXAMPLES_DIR = "examples" # مجلد جديد للأمثلة
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

MODELS_SPECIFIC_CONFIGS = {
    "interfuser_baseline": { "rgb_backbone_name": "r50", "embed_dim": 256, "direct_concat": True },
    "interfuser_lightweight": { "rgb_backbone_name": "r26", "embed_dim": 128, "enc_depth": 4, "dec_depth": 4, "direct_concat": True }
}

def find_available_models():
    if not os.path.isdir(WEIGHTS_DIR): return []
    return [f.replace(".pth", "") for f in os.listdir(WEIGHTS_DIR) if f.endswith(".pth")]

# ==============================================================================
#           2. الدوال الأساسية (load_model, run_single_frame)
# ==============================================================================
def load_model(model_name: str):
    """
    (لا تغيير في هذه الدالة)
    تبني وتحمل النموذج المختار وتُرجعه ككائن.
    """
    if not model_name or "لم يتم" in model_name:
        return None, "الرجاء اختيار نموذج صالح."

    weights_path = os.path.join(WEIGHTS_DIR, f"{model_name}.pth")
    print(f"Building model: '{model_name}'")
    
    model_config = MODELS_SPECIFIC_CONFIGS.get(model_name, {})
    model = build_interfuser_model(model_config)
    
    if not os.path.exists(weights_path):
        gr.Warning(f"ملف الأوزان '{weights_path}' غير موجود. النموذج سيعمل بأوزان عشوائية.")
    else:
        try:
            state_dic = torch.load(weights_path, map_location=device, weights_only=True)
            model.load_state_dict(state_dic)
            print(f"تم تحميل أوزان النموذج '{model_name}' بنجاح.")
        except Exception as e:
            gr.Warning(f"فشل تحميل الأوزان للنموذج '{model_name}': {e}.")
            
    model.to(device)
    model.eval()
    
    return model, f"تم تحميل نموذج: {model_name}"


def run_single_frame(
    model_from_state, # المدخل من gr.State
    rgb_image_path,
    rgb_left_image_path,
    rgb_right_image_path,
    rgb_center_image_path,
    lidar_image_path,
    measurements_path,
    target_point_list
):
    """
    (تم تعديل هذه الدالة)
    تعالج إطارًا واحدًا، وتقوم بتحميل النموذج الافتراضي إذا لزم الأمر لجلسات الـ API.
    """
    # --- تعديل للتعامل مع جلسات الـ API ---
    # إذا كانت هذه جلسة API جديدة (model_state فارغ)، قم بتحميل النموذج الافتراضي
    if model_from_state is None:
        print("API session detected or model not loaded. Loading default model...")
        available_models = find_available_models()
        if not available_models:
            raise gr.Error("لا توجد نماذج متاحة للتحميل في مجلد 'model/weights'.")
        
        default_model_name = available_models[0]
        model_to_use, _ = load_model(default_model_name)
    else:
        # إذا كان النموذج محملًا بالفعل (من جلسة متصفح)، استخدمه مباشرة
        model_to_use = model_from_state
    
    if model_to_use is None:
        raise gr.Error("فشل تحميل النموذج. تحقق من السجلات (Logs) في Hugging Face Space.")
    # --- نهاية التعديل ---

    try:
        # --- 1. قراءة ومعالجة المدخلات ---
        if not (rgb_image_path and measurements_path):
            raise gr.Error("الرجاء توفير الصورة الأمامية وملف القياسات على الأقل.")
        
        rgb_image_pil = Image.open(rgb_image_path).convert("RGB")
        rgb_left_pil = Image.open(rgb_left_image_path).convert("RGB") if rgb_left_image_path else rgb_image_pil
        rgb_right_pil = Image.open(rgb_right_image_path).convert("RGB") if rgb_right_image_path else rgb_image_pil
        rgb_center_pil = Image.open(rgb_center_image_path).convert("RGB") if rgb_center_image_path else rgb_image_pil

        front_tensor = transform(rgb_image_pil).unsqueeze(0).to(device)
        left_tensor = transform(rgb_left_pil).unsqueeze(0).to(device)
        right_tensor = transform(rgb_right_pil).unsqueeze(0).to(device)
        center_tensor = transform(rgb_center_pil).unsqueeze(0).to(device)

        if lidar_image_path:
            lidar_array = np.load(lidar_image_path)
            if lidar_array.max() > 0: lidar_array = (lidar_array / lidar_array.max()) * 255.0
            lidar_pil = Image.fromarray(lidar_array.astype(np.uint8)).convert('RGB')
        else:
            lidar_pil = Image.fromarray(np.zeros((112, 112, 3), dtype=np.uint8))
        lidar_tensor = lidar_transform(lidar_pil).unsqueeze(0).to(device)

        with open(measurements_path, 'r') as f: m_dict = json.load(f)

        measurements_tensor = torch.tensor([[
            m_dict.get('x',0.0), m_dict.get('y',0.0), m_dict.get('theta',0.0), m_dict.get('speed',5.0),
            m_dict.get('steer',0.0), m_dict.get('throttle',0.0), float(m_dict.get('brake',0.0)),
            m_dict.get('command',2.0), float(m_dict.get('is_junction',0.0)), float(m_dict.get('should_brake',0.0))
        ]], dtype=torch.float32).to(device)
        
        target_point_tensor = torch.tensor([target_point_list], dtype=torch.float32).to(device)
        
        inputs = {
            'rgb': front_tensor, 'rgb_left': left_tensor, 'rgb_right': right_tensor,
            'rgb_center': center_tensor, 'lidar': lidar_tensor,
            'measurements': measurements_tensor, 'target_point': target_point_tensor
        }

        # --- 2. تشغيل النموذج ---
        with torch.no_grad():
            outputs = model_to_use(inputs) # <-- استخدام model_to_use
            traffic, waypoints, is_junction, traffic_light, stop_sign, _ = outputs
        
        # --- 3. المعالجة اللاحقة والتصوّر ---
        speed, pos, theta = m_dict.get('speed',5.0), [m_dict.get('x',0.0), m_dict.get('y',0.0)], m_dict.get('theta',0.0)
        traffic_np, waypoints_np = traffic[0].detach().cpu().numpy().reshape(20,20,-1), waypoints[0].detach().cpu().numpy() * WAYPOINT_SCALE_FACTOR
        tracker, controller = Tracker(), InterfuserController(ControllerConfig())
        updated_traffic = tracker.update_and_predict(traffic_np.copy(), pos, theta, 0)
        steer, throttle, brake, metadata = controller.run_step(speed, waypoints_np, is_junction.sigmoid()[0,1].item(), traffic_light.sigmoid()[0,0].item(), stop_sign.sigmoid()[0,1].item(), updated_traffic)
        
        map_t0, counts_t0 = render(updated_traffic, t=0)
        map_t1, counts_t1 = render(updated_traffic, t=T1_FUTURE_TIME)
        map_t2, counts_t2 = render(updated_traffic, t=T2_FUTURE_TIME)
        wp_map = render_waypoints(waypoints_np)
        self_car_map = render_self_car(np.array([0,0]), [math.cos(0), math.sin(0)], [4.0, 2.0])
        map_t0 = cv2.add(cv2.add(map_t0, wp_map), self_car_map)
        map_t0 = cv2.resize(map_t0, (400, 400))
        map_t1 = cv2.add(ensure_rgb(map_t1), ensure_rgb(self_car_map)); map_t1 = cv2.resize(map_t1, (200, 200))
        map_t2 = cv2.add(ensure_rgb(map_t2), ensure_rgb(self_car_map)); map_t2 = cv2.resize(map_t2, (200, 200))
        display = DisplayInterface()
        light_state, stop_sign_state = "Red" if traffic_light.sigmoid()[0,0].item() > 0.5 else "Green", "Yes" if stop_sign.sigmoid()[0,1].item() > 0.5 else "No"
        interface_data = {'camera_view': np.array(rgb_image_pil),'map_t0': map_t0,'map_t1': map_t1,'map_t2': map_t2,
            'text_info': {'Control': f"S:{steer:.2f} T:{throttle:.2f} B:{int(brake)}",'Light': f"L: {light_state}",'Stop': f"St: {stop_sign_state}"},
            'object_counts': {'t0': counts_t0,'t1': counts_t1,'t2': counts_t2}}
        dashboard_image = display.run_interface(interface_data)
        
        # --- 4. تجهيز المخرجات ---
        result_dict = {"predicted_waypoints": waypoints_np.tolist(), "control_commands": {"steer": steer,"throttle": throttle,"brake": bool(brake)},
            "perception": {"traffic_light_status": light_state,"stop_sign_detected": (stop_sign_state == "Yes"),"is_at_junction_prob": round(is_junction.sigmoid()[0,1].item(), 3)},
            "metadata": {"speed_info": metadata[0],"perception_info": metadata[1],"stop_info": metadata[2],"safe_distance": metadata[3]}}
        
        return Image.fromarray(dashboard_image), result_dict
    except Exception as e:
        print(traceback.format_exc())
        raise gr.Error(f"حدث خطأ أثناء معالجة الإطار: {e}")


# ==============================================================================
#           4. تعريف واجهة Gradio المحسّنة (مع الإصلاح)
# ==============================================================================
available_models = find_available_models()

with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue", secondary_hue="sky"), css=".gradio-container {max-width: 95% !important;}") as demo:
    model_state = gr.State(value=None)
    
    gr.Markdown("# 🚗 محاكاة القيادة الذاتية باستخدام Interfuser")
    gr.Markdown("مرحباً بك في واجهة اختبار نموذج Interfuser. اتبع الخطوات أدناه لتشغيل المحاكاة على إطار واحد.")

    with gr.Row():
        # -- العمود الأيسر: الإعدادات والمدخلات --
        with gr.Column(scale=1):
            # --- الخطوة 1: اختيار النموذج ---
            with gr.Group():
                gr.Markdown("## ⚙️ الخطوة 1: اختر النموذج")
                with gr.Row():
                    model_selector = gr.Dropdown(
                        label="النماذج المتاحة",
                        choices=available_models,
                        value=available_models[0] if available_models else "لم يتم العثور على نماذج"
                    )
                    status_textbox = gr.Textbox(label="حالة النموذج", interactive=False)
            
            # --- الخطوة 2: رفع ملفات السيناريو ---
            with gr.Group():
                gr.Markdown("## 🗂️ الخطوة 2: ارفع ملفات السيناريو")
                
                with gr.Group():
                    gr.Markdown("**(مطلوب)**")
                    api_rgb_image_path = gr.File(label="صورة الكاميرا الأمامية (RGB)", type="filepath")
                    api_measurements_path = gr.File(label="ملف القياسات (JSON)", type="filepath")

                with gr.Accordion("📷 مدخلات اختيارية (كاميرات ومستشعرات إضافية)", open=False):
                    api_rgb_left_image_path = gr.File(label="كاميرا اليسار (RGB)", type="filepath")
                    api_rgb_right_image_path = gr.File(label="كاميرا اليمين (RGB)", type="filepath")
                    api_rgb_center_image_path = gr.File(label="كاميرا الوسط (RGB)", type="filepath")
                    api_lidar_image_path = gr.File(label="بيانات الليدار (NPY)", type="filepath")
                
                api_target_point_list = gr.JSON(label="📍 النقطة المستهدفة (x, y)", value=[0.0, 100.0])
                
                api_run_button = gr.Button("🚀 شغل المحاكاة", variant="primary", scale=2)

            # --- أمثلة جاهزة ---
            with gr.Group():
                gr.Markdown("### ✨ أمثلة جاهزة")
                gr.Markdown("انقر على مثال لتعبئة الحقول تلقائياً (يتطلب وجود مجلد `examples` بنفس بنية البيانات).")
                gr.Examples(
                    examples=[
                        [os.path.join(EXAMPLES_DIR, "sample1", "rgb.jpg"), os.path.join(EXAMPLES_DIR, "sample1", "measurements.json")],
                        [os.path.join(EXAMPLES_DIR, "sample2", "rgb.jpg"), os.path.join(EXAMPLES_DIR, "sample2", "measurements.json")]
                    ],
                    # يجب أن تتطابق المدخلات مع الحقول المطلوبة في الأمثلة
                    inputs=[api_rgb_image_path, api_measurements_path],
                    label="اختر سيناريو اختبار"
                )

        # -- العمود الأيمن: المخرجات --
        with gr.Column(scale=2):
            with gr.Group():
                gr.Markdown("## 📊 الخطوة 3: شاهد النتائج")
                api_output_image = gr.Image(label="لوحة التحكم المرئية (Dashboard)", type="pil", interactive=False)
                with gr.Accordion("عرض نتائج JSON التفصيلية", open=False):
                    api_output_json = gr.JSON(label="النتائج المهيكلة (JSON)")

    # --- ربط منطق الواجهة ---
    if available_models:
        demo.load(fn=load_model, inputs=model_selector, outputs=[model_state, status_textbox])
    
    model_selector.change(fn=load_model, inputs=model_selector, outputs=[model_state, status_textbox])
    
    api_run_button.click(
        fn=run_single_frame,
        inputs=[model_state, api_rgb_image_path, api_rgb_left_image_path, api_rgb_right_image_path, 
                api_rgb_center_image_path, api_lidar_image_path, api_measurements_path, api_target_point_list],
        outputs=[api_output_image, api_output_json],
        api_name="run_single_frame"
    )

# ==============================================================================
#           5. تشغيل التطبيق
# ==============================================================================
if __name__ == "__main__":
    if not available_models:
        print("تحذير: لم يتم العثور على أي ملفات نماذج (.pth) في مجلد 'model/weights'.")
    demo.queue().launch(debug=True, share=True)