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

ШІ-видалення фону з зображень

Завдання відділення об'єкту від фону ділиться на два класи: грубе видалення прямокутного фону для e-commerce та точне альфа-матування (alpha matting) для волосся, хутра, напівпрозорих об'єктів. Ці завдання технічно відрізняються і вирішуються різними інструментами.

Grounded-SAM2 — сучасний стандарт

Segment Anything Model 2 (Meta, 2024) у поєднанні з Grounding DINO забезпечує найсучаснішу якість для більшості завдань. 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

Альфа-матування для складних границь

Волосся, хутро, тонкі гілки — SAM2 дає грубу маску за bbox, границі стають пікселеутворюючими. Для цих випадків поверх SAM-маски застосовується альфа-матування:

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)

Пакетна обробка для E-commerce

Інтернет-магазини обробляють тисячі фото товарів. Оптимальна схема: 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

Порівняння інструментів

Часові рамки