Update app.py

app.py

@@ -1,45 +1,47 @@
 import torch
-from torchvision import transforms
 from PIL import Image
-import gradio as gr
-from transformers import AutoModelForImageClassification, AutoFeatureExtractor
 import numpy as np
+import gradio as gr
+from transformers import AutoFeatureExtractor
+from huggingface_hub import snapshot_download
+import sys
+import os
 
-# تحميل النموذج مع السماح بتشغيل الكود المخصص
-birefnet = AutoModelForImageClassification.from_pretrained(
-    "briaai/RMBG-1.4", trust_remote_code=True
-)
-birefnet.to("cpu")  # ✅ التشغيل على CPU
+# تحميل ملفات النموذج بالكامل
+model_path = snapshot_download(repo_id="briaai/RMBG-1.4", trust_remote_code=True)
 
-# تحميل محول الميزات (feature extractor)
-extractor = AutoFeatureExtractor.from_pretrained("briaai/RMBG-1.4")
+# إضافة مجلد النموذج إلى المسار
+sys.path.insert(0, model_path)
 
-# تحويل الصورة لإدخال النموذج
-def transform_image(image):
-    inputs = extractor(images=image, return_tensors="pt")
-    return inputs["pixel_values"][0]
+# استيراد الكلاس المخصص للنموذج
+from model import RMBGModel
+
+# تهيئة النموذج
+model = RMBGModel.from_pretrained(model_path, trust_remote_code=True)
+model.eval()
+
+# تحميل الـ feature extractor
+extractor = AutoFeatureExtractor.from_pretrained("briaai/RMBG-1.4")
 
-# دالة إزالة الخلفية
 def process(image):
-    input_images = transform_image(image).unsqueeze(0).to("cpu")
+    inputs = extractor(images=image, return_tensors="pt")
     with torch.no_grad():
-        output = birefnet(input_images).logits.squeeze(0)[0]
-    mask = torch.sigmoid(output).cpu().numpy()
-    mask = (mask * 255).astype(np.uint8)
-    mask = Image.fromarray(mask).resize(image.size)
+        result = model(**inputs)
+    alpha = result["output"].squeeze().numpy()
+    alpha = (alpha * 255).astype(np.uint8)
+    alpha = Image.fromarray(alpha).resize(image.size)
 
     image = image.convert("RGBA")
-    mask = mask.convert("L")
-    image.putalpha(mask)
+    alpha = alpha.convert("L")
+    image.putalpha(alpha)
     return image
 
-# واجهة Gradio
 demo = gr.Interface(
     fn=process,
     inputs=gr.Image(type="pil"),
     outputs=gr.Image(type="pil"),
-    title="إزالة خلفية الصور باستخدام BiRefNet (CPU)",
-    description="ارفع صورة وسيتم إزالة الخلفية تلقائيًا باستخدام نموذج BiRefNet على وحدة المعالجة المركزية فقط."
+    title="إزالة خلفية الصور باستخدام RMBG",
+    description="ارفع صورة وسيتم إزالة الخلفية تلقائيًا باستخدام نموذج RMBG من briaai."
 )
 
 demo.launch()