Add initial implementation of image segmentation app with SAM2 model and necessary files

.gitignore

@@ -0,0 +1,207 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[codz]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py.cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# UV
+#   Similar to Pipfile.lock, it is generally recommended to include uv.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#uv.lock
+
+# poetry
+#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
+#poetry.lock
+#poetry.toml
+
+# pdm
+#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
+#   pdm recommends including project-wide configuration in pdm.toml, but excluding .pdm-python.
+#   https://pdm-project.org/en/latest/usage/project/#working-with-version-control
+#pdm.lock
+#pdm.toml
+.pdm-python
+.pdm-build/
+
+# pixi
+#   Similar to Pipfile.lock, it is generally recommended to include pixi.lock in version control.
+#pixi.lock
+#   Pixi creates a virtual environment in the .pixi directory, just like venv module creates one
+#   in the .venv directory. It is recommended not to include this directory in version control.
+.pixi
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.envrc
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# PyCharm
+#  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
+#  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
+#  and can be added to the global gitignore or merged into this file.  For a more nuclear
+#  option (not recommended) you can uncomment the following to ignore the entire idea folder.
+#.idea/
+
+# Abstra
+# Abstra is an AI-powered process automation framework.
+# Ignore directories containing user credentials, local state, and settings.
+# Learn more at https://abstra.io/docs
+.abstra/
+
+# Visual Studio Code
+#  Visual Studio Code specific template is maintained in a separate VisualStudioCode.gitignore 
+#  that can be found at https://github.com/github/gitignore/blob/main/Global/VisualStudioCode.gitignore
+#  and can be added to the global gitignore or merged into this file. However, if you prefer, 
+#  you could uncomment the following to ignore the entire vscode folder
+# .vscode/
+
+# Ruff stuff:
+.ruff_cache/
+
+# PyPI configuration file
+.pypirc
+
+# Cursor
+#  Cursor is an AI-powered code editor. `.cursorignore` specifies files/directories to
+#  exclude from AI features like autocomplete and code analysis. Recommended for sensitive data
+#  refer to https://docs.cursor.com/context/ignore-files
+.cursorignore
+.cursorindexingignore
+
+# Marimo
+marimo/_static/
+marimo/_lsp/
+__marimo__/

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 Kenny Santanu
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

app.py

@@ -0,0 +1,172 @@
+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 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)
+
+            # 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)])
+
+            # 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."]
+
+            # 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)
+
+            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}"]
+
+    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", 
+                    type="pil",
+                    height=400
+                )
+                result_mask = gr.Image(
+                    label="SAM2 Mask",
+                    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
+
+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
+    )
+
+# Run the app if this script is executed directly
+if __name__ == "__main__":
+    main()

requirements.txt

@@ -0,0 +1,3 @@
+gradio>=5.37.0
+opencv-python>=4.12.0.88
+ultralytics>=8.3.167