Реализация AI-удаления фона изображения

AI-удаление фона с изображений

Задача отделения объекта от фона делится на два класса: грубое удаление прямоугольного фона для ecommerce и точное альфа-матирование (alpha matting) для волос, меха, полупрозрачных объектов. Эти задачи технически различаются и решаются разными инструментами.

Grounded-SAM2 — современный стандарт

Segment Anything Model 2 (Meta, 2024) в связке с Grounding DINO даёт state-of-the-art качество для большинства задач. SAM2 — promptable: передаём точку или bbox, получаем маску.

import torch
import numpy as np
from PIL import Image
from sam2.build_sam import build_sam2
from sam2.sam2_image_predictor import SAM2ImagePredictor
from groundingdino.util.inference import load_model, predict

def remove_background_grounded_sam2(
    image_path: str,
    text_prompt: str = 'product',   # что вырезать
    box_threshold: float = 0.3,
    text_threshold: float = 0.25,
    output_path: str = None
) -> Image.Image:
    image = Image.open(image_path).convert('RGB')
    image_np = np.array(image)

    # Grounding DINO → bbox объекта по тексту
    gdino_model = load_model(
        'groundingdino/config/GroundingDINO_SwinT_OGC.py',
        'weights/groundingdino_swint_ogc.pth'
    )
    boxes, _, _ = predict(
        model=gdino_model,
        image=image_np,
        caption=text_prompt,
        box_threshold=box_threshold,
        text_threshold=text_threshold
    )

    if len(boxes) == 0:
        raise ValueError(f'Object "{text_prompt}" not found')

    # SAM2 → точная маска по bbox
    sam2 = build_sam2(
        'sam2_hiera_large.yaml',
        'weights/sam2_hiera_large.pt',
        device='cuda'
    )
    predictor = SAM2ImagePredictor(sam2)
    predictor.set_image(image_np)

    # Берём самый уверенный bbox
    best_box = boxes[0].numpy() * np.array([
        image_np.shape[1], image_np.shape[0],
        image_np.shape[1], image_np.shape[0]
    ])

    masks, scores, _ = predictor.predict(
        box=best_box,
        multimask_output=True
    )
    best_mask = masks[np.argmax(scores)]  # (H, W) bool

    # Применяем маску → RGBA
    result_rgba = np.dstack([image_np, best_mask.astype(np.uint8) * 255])
    result = Image.fromarray(result_rgba, 'RGBA')

    if output_path:
        result.save(output_path, 'PNG')

    return result

Alpha matting для сложных краёв

Волосы, мех, тонкие ветки — SAM2 даёт грубую маску по bbox, края получаются пиксельными. Для этих случаев поверх SAM-маски применяется alpha matting:

from pymatting import estimate_alpha_cf, estimate_foreground_ml
import cv2

def refine_mask_with_matting(
    image: np.ndarray,    # (H, W, 3) RGB
    rough_mask: np.ndarray,  # (H, W) bool от SAM
    erosion_px: int = 10,    # размер «неопределённой» зоны
    dilation_px: int = 10
) -> np.ndarray:
    """
    Уточнение маски через closed-form matting.
    Trimap: definite_fg=255, definite_bg=0, uncertain=128
    """
    kernel = np.ones((erosion_px, erosion_px), np.uint8)

    fg_mask = cv2.erode(
        rough_mask.astype(np.uint8) * 255, kernel
    )
    bg_mask = cv2.dilate(
        rough_mask.astype(np.uint8) * 255, kernel
    )

    trimap = np.full(rough_mask.shape, 128, dtype=np.uint8)
    trimap[fg_mask > 0] = 255
    trimap[bg_mask == 0] = 0

    # Closed-form matting (Levin et al.)
    image_float = image.astype(np.float64) / 255.0
    trimap_float = trimap.astype(np.float64) / 255.0

    alpha = estimate_alpha_cf(image_float, trimap_float)
    # alpha ∈ [0, 1], значения на краях — дробные (полупрозрачность)

    return (alpha * 255).astype(np.uint8)

Батчевая обработка для ecommerce

Интернет-магазины обрабатывают тысячи товарных фото. Оптимальная схема: REMBG (библиотека на базе U2-Net/IS-Net) для быстрого батч-инференса, SAM2 — для спорных случаев.

from rembg import remove, new_session
from PIL import Image
from pathlib import Path
import concurrent.futures

def batch_remove_background(
    input_dir: str,
    output_dir: str,
    model_name: str = 'isnet-general-use',  # лучше для товаров
    max_workers: int = 4
) -> dict:
    """
    Модели rembg: u2net, u2netp (быстрее, хуже),
    isnet-general-use (качество), isnet-anime (аниме).
    """
    session = new_session(model_name)
    input_paths = list(Path(input_dir).glob('*.{jpg,jpeg,png,webp}'))
    results = {'success': 0, 'failed': 0, 'errors': []}

    def process_one(img_path: Path) -> bool:
        try:
            with open(img_path, 'rb') as f:
                input_data = f.read()
            output_data = remove(input_data, session=session)
            out_path = Path(output_dir) / (img_path.stem + '.png')
            with open(out_path, 'wb') as f:
                f.write(output_data)
            return True
        except Exception as e:
            results['errors'].append(str(e))
            return False

    with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as ex:
        futures = {ex.submit(process_one, p): p for p in input_paths}
        for fut in concurrent.futures.as_completed(futures):
            if fut.result():
                results['success'] += 1
            else:
                results['failed'] += 1

    return results

Сравнение инструментов

Сроки