Object Segmentation and Movement Tool

This project is a Python-based tool for object segmentation and movement using state-of-the-art deep learning models. The script leverages Grounding DINO, SAM (Segment Anything Model), and ControlNet models for segmentation and object inpainting.

Methodology

1. Overview:

   • Grounding DINO is used for text-based object detection.
   • SAM (Segment Anything Model) segments the object from the image using bounding boxes generated by Grounding DINO.
   • ControlNet with Stable Diffusion inpainting handles the inpainting and relocation tasks.

2. Text-based Object Detection and Segmentation:

   • Input an image and a text prompt describing the object.
   • Use Grounding DINO to identify bounding boxes matching the prompt.
   • Generate masks using SAM for precise segmentation.

3. Object Relocation and Inpainting:

   • Mask the detected object and inpaint the original area with photorealistic background using Stable Diffusion.
   • Create a shifted mask to define the new location.
   • Use ControlNet-guided inpainting to seamlessly integrate the object into its new position.

4. Output Generation:

   • Save intermediate results (bounding boxes, masks, cropped objects).
   • Produce final images showing segmentation or relocated objects.

Requirements

Libraries

• Python 3.11+
• PyTorch
• NumPy
• OpenCV
• Pillow
• Diffusers
• Transformers
• Matplotlib

Install the required libraries using:

pip install -r requirements.txt

Models

The script uses pretrained models:

• Grounding DINO: IDEA-Research/grounding-dino-base
• SAM: facebook/sam-vit-huge
• ControlNet Inpainting: lllyasviel/control_v11p_sd15_inpaint
• Stable Diffusion: runwayml/stable-diffusion-inpainting

We use the implementation of these models present on Diffusers and Transformers and are automatically downloaded when the script is executed

Usage

Command-line Arguments

python run.py --image <image_path> --class <object_class> [--output <output_path>] [--x <x_shift>] [--y <y_shift>]

Arguments:

• --image: Path to the input image.
• --class: Text description of the object to segment.
• --output: Path to save the output image (optional).
• --x: Horizontal shift in pixels for moving the object (default: 0).
• --y: Vertical shift in pixels for moving the object (default: 0).

Examples

Segmentation Only

To segment an object:

python run.py --image input.jpg --class "chair" --output segmented_output.png

Object Movement

To move an object:

python run.py --image input.jpg --class "chair" --output moved_output.png --x 50 --y 100

Output

1. Segmentation Visualization:

   • A semi-transparent red overlay is applied to the segmented object.
   • Saved in ./sam_mask/ or the specified output path.

2. Object Movement:

   • Relocated object with photorealistic inpainting.
   • Saved in ./outputs/ or the specified output path.

Directory Structure

The script creates and uses the following directories for intermediate results:

• ./dino_bbox/: Bounding box images from Grounding DINO.
• ./sam_mask/: Masks generated by SAM.
• ./extracted_object/: Cropped object images.
• ./inpainted_background/: Background inpainting results (first inpainting without shifting of the object)
• ./outputs/: Final results after object movement.

Notes

• Ensure CUDA is available for GPU acceleration; otherwise, computations will default to the CPU.
• The script handles cases where the described object is not detected by using pre-defined bounding boxes for specific objects.

License

This project is licensed under the MIT License. See the LICENSE file for details.

Acknowledgments

• Hugging Face for the Transformers library.
• Diffusers for inpainting and ControlNet pipelines.
• Open source contributors for Grounding DINO and SAM.