from controlnet_aux import LineartDetector
from PIL import Image
import cv2
import numpy as np

def extract_petroglyph_with_ai(image_path):
    print("جاري تحميل نموذج الذكاء الاصطناعي (قد يستغرق بضع ثوانٍ في المرة الأولى)...")
    
    # 1. تحميل النموذج المتخصص في استخراج الرسومات الخطية
    # سيقوم الكود بتحميل الأوزان (حوالي 150 ميجابايت) في أول مرة تشغيل فقط
    processor = LineartDetector.from_pretrained("lllyasviel/Annotators")

    # 2. قراءة الصورة الأصلية
    try:
        image = Image.open(image_path)
    except FileNotFoundError:
        print("خطأ: تأكد من صحة مسار واسم الصورة.")
        return

    print("جاري تحليل الصخرة واستخراج النقوش...")
    
    # 3. تمرير الصورة للنموذج (الناتج الافتراضي: خطوط بيضاء على خلفية سوداء)
    # coarse=True تساعد أحياناً في التقاط الخطوط العريضة وتجاهل الخدوش الرفيعة
    processed_image = processor(image, coarse=True)

    # 4. تحويل النتيجة إلى مصفوفة للتمكن من تعديل ألوانها
    open_cv_image = np.array(processed_image)

    # 5. عكس الألوان: تحويل الخلفية إلى أبيض ناصع، والنقش إلى أسود
    white_paper_result = cv2.bitwise_not(open_cv_image)

    # 6. حفظ النتيجة كصورة نهائية
    output_filename = "ai_extracted_petroglyph.png"
    cv2.imwrite(output_filename, white_paper_result)
    
    print(f"تمت العملية بنجاح! تم حفظ النتيجة كـ: {output_filename}")

    # عرض النتيجة على الشاشة
    cv2.imshow("AI Result - Press any key to close", white_paper_result)
    cv2.waitKey(0)
    cv2.destroyAllWindows()

# ضع مسار صورة الصخرة هنا
extract_petroglyph_with_ai("rock.jpg")