Update app.py

app.py

@@ -1,172 +1,125 @@
 import gradio as gr
 import numpy as np
-import cv2
 from PIL import Image
-from ultralytics import SAM
 
-class ImageSegmentationApp:
-    def __init__(self) -> None:
-        """Initialize the segmentation app and load the SAM2 model with fallback."""
-        try:
-            # Attempt to load the SAM2 model weights
-            self.model = SAM("sam2.1_t.pt")
-            self.model_available = True  # Model loaded successfully
-        except Exception as e:
-            # If loading fails, set model as unavailable and print error
-            print(f"Failed to load SAM2 model: {e}")
-            self.model = None
-            self.model_available = False
+def get_mask_bbox(mask: Image.Image):
+    """احصل على إطار القناع (Bounding Box) كـ (xmin, ymin, xmax, ymax)"""
+    mask_np = np.array(mask)
+    if len(mask_np.shape) == 3:
+        mask_np = mask_np[..., -1] # قناة ألفا
+    ys, xs = np.where(mask_np > 0)
+    if ys.size == 0 or xs.size == 0:
+        return None
+    return int(xs.min()), int(ys.min()), int(xs.max()), int(ys.max())
 
-    def process_segmentation(
-        self, 
-        image_editor: dict, 
-        replacement_image: Image.Image
-    ) -> list[Image.Image | None] | None:
-        """
-        Process the segmentation and replacement using the drawn mask and SAM2 model.
-        Returns [drawn_mask, sam_mask, result_image, markdown_message].
-        """
-        # Check if both images are provided
-        if image_editor["background"] is None or replacement_image is None:
-            return [None, None, None, "**❌ Error:** Please upload both images."]
-        try:
-            # Extract the original image and the user-drawn mask
-            original_image = image_editor["background"]
-            drawn_mask = image_editor["layers"][0]
-            
-            # Use the alpha channel of the mask as the binary mask
-            drawn_mask = drawn_mask.split()[-1]
-            drawn_mask_np = np.array(drawn_mask)
+def crop_masked_object(img: Image.Image, mask: Image.Image):
+    """اقطع المنطقة المحددة من الصورة حسب القناع وأرجعها"""
+    bbox = get_mask_bbox(mask)
+    if bbox is None:
+        return None, None
+    xmin, ymin, xmax, ymax = bbox
+    # قص الصورة والقناع للمنطقة المحددة فقط
+    cropped_img = img.crop((xmin, ymin, xmax + 1, ymax + 1))
+    cropped_mask = mask.crop((xmin, ymin, xmax + 1, ymax + 1))
+    # طبّق القناع على الجزء المقصوص (شفافية)
+    cropped_img.putalpha(cropped_mask.split()[-1])
+    return cropped_img, bbox
 
-            # Find contours in the mask to determine segmentation points
-            points = []
-            contours, _ = cv2.findContours(drawn_mask_np, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
-            for contour in contours:
-                M = cv2.moments(contour)
-                if M["m00"] != 0:
-                    # Use centroid of contour as a point
-                    cx = float(M["m10"] / M["m00"])
-                    cy = float(M["m01"] / M["m00"])
-                    points.append([cx, cy])
-                else:
-                    # Fallback: use the first point if the area is zero
-                    x, y = contour[0][0]
-                    points.append([float(x), float(y)])
+def paste_object_on_mask(
+    base_img: Image.Image,
+    base_mask: Image.Image,
+    obj_img: Image.Image,
+    obj_bbox,
+):
+    """الصق الجزء في منطقة القناع على الصورة الأساسية مع التحجيم والمحاذاة"""
+    base_bbox = get_mask_bbox(base_mask)
+    if base_bbox is None or obj_bbox is None:
+        return None
+    bxmin, bymin, bxmax, bymax = base_bbox
+    ow, oh = obj_img.size
+    # حجم منطقة اللصق في الصورة الأصلية
+    bw, bh = bxmax - bxmin + 1, bymax - bymin + 1
+    # غيّر حجم الجزء المقصوص ليناسب منطقة اللصق
+    obj_img_resized = obj_img.resize((bw, bh), resample=Image.LANCZOS)
+    # لصق الجزء على الصورة الأصلية في مكان القناع
+    result = base_img.convert("RGBA")
+    mask = base_mask.split()[-1]
+    empty = Image.new("RGBA", result.size, (0, 0, 0, 0))
+    empty.paste(obj_img_resized, (bxmin, bymin), obj_img_resized)
+    # الدمج: الجزء الجديد مكان القناع فقط
+    result = Image.composite(empty, result, mask)
+    return result
 
-            # If no points are found, return original image and a message indicating no mask was drawn         
-            if not points:
-                return [None, None, original_image, "**❌ Error:** No mask drawn. Please draw a mask on the original image."]
+def process_transfer(
+    image1_editor, # dict: {"background": Image, "layers": [Image,...]}
+    image2_editor  # dict: {"background": Image, "layers": [Image,...]}
+):
+    # التحقق من المدخلات
+    if not image1_editor or not image2_editor:
+        return None, None, None, "**❌ يرجى رفع الصورتين وتحديد الأقنعة."
+    if (image1_editor.get("background") is None or
+        not image1_editor.get("layers") or
+        image2_editor.get("background") is None or
+        not image2_editor.get("layers")
+    ):
+        return None, None, None, "**❌ يجب رسم قناع على كلتا الصورتين."
+    base_img = image1_editor["background"].convert("RGBA")
+    base_mask = image1_editor["layers"][0].convert("RGBA")
+    obj_img = image2_editor["background"].convert("RGBA")
+    obj_mask = image2_editor["layers"][0].convert("RGBA")
 
-            # If the SAM2 model is unavailable, use the drawn mask directly
-            if not self.model_available or not self.model:
-                sam_mask = drawn_mask
-                model_message = "**⚠️ Warning:** SAM2 model unavailable, using drawn mask as mask."
-            else:
-                # Run the SAM2 model to refine the mask
-                results = self.model(
-                    source=original_image,
-                    points=[points],
-                )
-                # Extract the mask from the model output
-                result_numpy_arr = results[0].masks.data.numpy()
-                sam_mask_arr = np.squeeze(result_numpy_arr)
-                sam_mask_arr = (sam_mask_arr * 255).astype(np.uint8)  # Convert bool to uint8
-                sam_mask = Image.fromarray(sam_mask_arr)
-                model_message = "**✅ Success:** Segmentation completed with SAM2."
-
-            # Resize the replacement image to match the original image size
-            replacement_image = replacement_image.resize(original_image.size)
-            # Composite the replacement image onto the original using the mask
-            result_image = Image.composite(replacement_image, original_image, sam_mask)
+    # قص الجزء المحدد من الصورة الثانية
+    cropped_obj, obj_bbox = crop_masked_object(obj_img, obj_mask)
+    if cropped_obj is None:
+        return None, None, None, "**❌ لم يتم تحديد أي كائن في الصورة الثانية."
+    # لصق الجزء في منطقة القناع بالصورة الأولى
+    result_img = paste_object_on_mask(base_img, base_mask, cropped_obj, obj_bbox)
+    if result_img is None:
+        return None, None, None, "**❌ لم يتم تحديد منطقة مناسبة للصق في الصورة الأولى."
 
-            return [drawn_mask, sam_mask, result_image, model_message]
-        except Exception as e:
-            # Catch and report any errors during segmentation
-            print(f"Segmentation error: {e}")
-            return [None, None, None, f"**❌ Error:** Segmentation error: {e}"]
+    return base_mask, obj_mask, result_img, "**✅ تم نقل الكائن بنجاح!"
 
-    def create_interface(self) -> gr.Blocks:
-        """Create and return the Gradio interface"""
-        with gr.Blocks(title="SAM2 Image Segmentation & Replacement", theme=gr.themes.Soft(), css=".center-status-message {text-align: center;}") as demo:
-            # App title and instructions
-            gr.Markdown(
-                f"""
-                # 🎨 SAM2 Image Segmentation & Replacement
-                
-                Upload an original image and a replacement image, then draw a rough mask on the original image.
-                
-                **Instructions:**
-                1. Upload your original image
-                2. Upload your replacement image  
-                3. Draw a mask on the original image by painting over the area you want to replace
-                4. Click "Process Segmentation" to see the result
-                """
-            )
-            gr.Markdown("### 📸 Upload Images")
-            with gr.Row():
-                with gr.Column():
-                    # ImageMask for original image and mask drawing
-                    image_editor = gr.ImageMask(
-                        label="Original Image",
-                        type="pil",
-                        height=400
-                    )
-                with gr.Column():
-                    # Upload for replacement image
-                    replacement_image = gr.Image(
-                        label="Replacement Image",
-                        type="pil",
-                        height=400
-                    )
-            with gr.Row():
-                # Button to trigger segmentation
-                process_btn = gr.Button("🚀 Process Segmentation", variant="primary", size="lg")
-            with gr.Row():
-                # Status message for feedback
-                status_message = gr.Markdown(value="", elem_id="status_message", elem_classes=["center-status-message"])
-            with gr.Row():
-                # Display the drawn mask, SAM2 mask, and result image
-                drawn_mask = gr.Image(
-                    label="Drawn Mask", 
+def create_interface():
+    with gr.Blocks(title="نقل كائن من صورة إلى أخرى مع القناع والمحاذاة") as demo:
+        gr.Markdown("""
+        # 🖼️ نقل كائن من صورة إلى أخرى مع القناع والمحاذاة
+        
+        1. ارفع الصورة الأولى (صورة الخلفية).
+        2. ارفع الصورة الثانية (الصورة التي بها الكائن).
+        3. ارسم قناعًا على الكائن في الصورة الثانية (صورة الاستبدال).
+        4. ارسم قناعًا في مكان اللصق في الصورة الأولى.
+        5. اضغط زر "نقل الكائن" وستظهر النتيجة!
+        """)
+        with gr.Row():
+            with gr.Column():
+                image1_editor = gr.ImageMask(
+                    label="الصورة الأولى (اختر مكان اللصق)",
                     type="pil",
                     height=400
                 )
-                result_mask = gr.Image(
-                    label="SAM2 Mask",
+            with gr.Column():
+                image2_editor = gr.ImageMask(
+                    label="الصورة الثانية (حدد الكائن المراد نقله)",
                     type="pil",
                     height=400
                 )
-                result_image = gr.Image(
-                    label="Result", 
-                    type="pil",
-                    height=400
-                )
-            with gr.Row():
-                # Display copywrite information
-                gr.Markdown(
-                    value="© 2025 Kenny Santanu. All rights reserved.",
-                    elem_classes=["center-status-message"]
-                )
-
-            # Connect button click to segmentation function
-            process_btn.click(
-                fn=self.process_segmentation,
-                inputs=[image_editor, replacement_image],
-                outputs=[drawn_mask, result_mask, result_image, status_message]
-            )
-        return demo
+        with gr.Row():
+            transfer_btn = gr.Button("🚀 نقل الكائن", variant="primary")
+        with gr.Row():
+            base_mask = gr.Image(label="قناع الصورة الأولى", type="pil", height=200)
+            obj_mask = gr.Image(label="قناع الصورة الثانية", type="pil", height=200)
+            result_img = gr.Image(label="النتيجة النهائية", type="pil", height=400)
+        status = gr.Markdown()
+        transfer_btn.click(
+            fn=process_transfer,
+            inputs=[image1_editor, image2_editor],
+            outputs=[base_mask, obj_mask, result_img, status]
+        )
+    return demo
 
-def main() -> None:
-    """Main function to run the application"""
-    # Instantiate the app
-    app = ImageSegmentationApp()
-    # Create the Gradio interface
-    demo = app.create_interface()
-    # Launch the interface (web server)
-    demo.launch(
-        show_api=False
-    )
+def main():
+    demo = create_interface()
+    demo.launch()
 
-# Run the app if this script is executed directly
 if __name__ == "__main__":
-    main()
+    main()