Segment Anything - Background Removal Tool

This project uses Meta's Segment Anything Model (SAM) to automatically generate masks and remove backgrounds from images. It was designed for batch processing large image datasets with cloud storage integration (Yandex Disk).

📋 Overview

The notebook NewMaskSelection.ipynb implements a complete pipeline for:

• Downloading images from Yandex Disk in batches
• Generating segmentation masks using SAM
• Selecting and refining the best mask for background removal
• Saving processed images with transparent backgrounds
• Uploading results back to cloud storage

🔧 Key Components

1. Model Setup

• Uses SAM (Segment Anything Model) with vit_l (Large ViT) architecture
• Pretrained weights: sam_vit_l_0b3195.pth
• Runs on CUDA-enabled GPUs for acceleration
• Employs SamAutomaticMaskGenerator for automatic mask generation

2. Core Functions

Image Preprocessing

• checkResolution(image, maxres) - Resizes images exceeding maximum resolution to prevent memory issues

Mask Generation & Selection

• sortMasksByArea(masks) - Sorts masks by area (largest first)
• sortMasksByPerimeter(masks) - Sorts masks by bounding box perimeter
• selectAndDeleteMask(image, masks, ...) - Main function that:
  • Selects the primary mask based on size criteria
  • Handles edge cases where SAM separates background into layers
  • Applies edge detection to preserve image boundaries
  • Generates final transparency mask

Mask Validation

• check_mask(mask, biggest_contour) - Validates mask covers sufficient image area
• checking_lines(mask, thickness, available) - Detects if mask has edge artifacts by checking border regions
• get_biggest_contour(contours) - Extracts largest contour for validation

Output Generation

• SaveAllMasks(image, masks, path, ext) - Saves top 8 masks for manual review
• make_all_calcs_and_save(...) - Complete processing pipeline for single image

3. Cloud Storage Integration (Yandex Disk)

Async Download

getBatchAsync(batch_id, client, batch_size=100, n_parallel_requests=10)

• Downloads batches of images from Yandex Disk
• Parallel downloads with configurable concurrency
• Creates local directory structure automatically

Async Upload

recursive_upload(from_dir, to_dir, client, n_parallel_requests=5)

• Recursively uploads processed results
• Creates remote directories as needed
• Parallel uploads for efficiency

📊 Workflow

┌─────────────────┐
│  Yandex Disk    │
│  (Source Images)│
└────────┬────────┘
         │ Download (async batch)
         ▼
┌─────────────────┐
│  Local Storage  │
│  E:/input/      │
└────────┬────────┘
         │ Process each image
         ▼
┌─────────────────┐
│  SAM Model      │
│  Generate Masks │
└────────┬────────┘
         │ Select & Refine
         ▼
┌─────────────────┐
│  Mask Selection │
│  & Validation   │
└────────┬────────┘
         │ Save Results
         ▼
┌─────────────────┐
│  Output Folders │
│  - res/         │ → Final images with transparent BG
│  - checking/    │ → Preview collages
│  - [name](id)/  │ → Individual mask variants
└────────┬────────┘
         │ Upload (async)
         ▼
┌─────────────────┐
│  Yandex Disk    │
│  (Results)      │
└─────────────────┘

📁 Output Structure

For each batch (e.g., batch_0/):

batch_0/
├── image1.png/
│   ├── image1.png_mask.png      # Primary mask
│   └── _(1).png, _(2).png...    # Alternative masks (top 8)
├── image2.png/
│   └── ...
├── checking/
│   ├── image1.png.png           # Collage: result | mask | original
│   └── ...
└── res/
    ├── image1.png               # Final transparent image
    └── ...

🎯 Mask Selection Logic

The notebook implements intelligent mask selection:

1. Primary Check: If largest mask covers ≥90% of image dimensions → use as foreground
2. Alternative: If not, combine all masks and invert (assumes background was segmented)
3. Edge Refinement:
   • Apply Gaussian blur to smooth mask edges
   • Check border regions (5-12px thickness)
   • If borders contain too much transparency → restore edge pixels
4. Validation: Verify mask contour spans full image height/width

📈 Performance

From execution logs:

• ~1500+ images processed in Google Colab initially
• Full reprocessing done locally for better reliability
• Processing time: ~2-5 seconds per batch (100 images)
• Total processing: ~5 minutes for 150+ batches

🛠️ Requirements

Python Libraries

torch (2.2.2+cu121)
opencv-python (cv2)
numpy
matplotlib
yadisk (Yandex Disk API)
asyncio

Hardware

• GPU: CUDA-enabled (tested with torch 2.2.2+cu121)
• RAM: Sufficient for large image batches
• Storage: Local directories for input/output caching

Model Files

• sam_vit_l_0b3195.pth - SAM ViT-Large weights

📝 Usage

Setup

1. Install dependencies: pip install torch opencv-python numpy matplotlib yadisk
2. Download SAM checkpoint: sam_vit_l_0b3195.pth
3. Configure paths in notebook:
   • input_path - Local output directory
   • input_dir - Local input directory
   • download_path - Yandex Disk source path
   • upload_path - Yandex Disk result path

Execution

Run notebook cells sequentially:

1. Import libraries
2. Define async cloud functions
3. Define mask processing functions
4. Load SAM model
5. Validate batches (check for errors/missing files)
6. Process new images
7. Upload results

🔍 Error Handling

The notebook tracks three error categories:

1. list_of_inv_bg - Images with invalid background masks (1432 found in sample)
2. list_of_missed_res_check - Masks generated but results missing (0 in sample)
3. list_of_missed_masks - Mask generation failed (0 in sample)

Common issues identified:

• .psd files or 0-byte files mislabeled as images
• SAM sometimes segments background instead of foreground
• Edge artifacts in masks require post-processing

📸 Example Results

Preview of background removal: https://disk.yandex.ru/d/L8SmEplUzZMJpg

🏗️ Development Notes

• Originally designed for Google Colab with dynamic Yandex Disk integration
• Moved to local execution for better stability and performance
• OpenCV errors led to discovery of corrupted input files
• Increased parallel request count for efficient small file uploads
• Multiple reprocessing runs to refine mask selection algorithm

📄 License

This project uses Meta's Segment Anything Model. See SAM license for usage terms.