difference_matting

diff_mat.py

#!/usr/bin/env python3
"""
Difference matting script to extract transparency from two images:
one rendered on a black background, one on a white background.

Usage: python stitch.py <black_bg_image> <white_bg_image> <output.png>
"""

import sys
import math
from PIL import Image
import numpy as np


def extract_alpha(black_bg_path: str, white_bg_path: str, output_path: str) -> None:
    """
    Extract transparency using difference matting technique.

    The theory: If a pixel is 100% opaque, it looks identical on both backgrounds.
    If it's 100% transparent, it shows the background color exactly.
    The distance between the two observations tells us the alpha value.
    """
    # Load images
    img_black = Image.open(black_bg_path).convert('RGB')
    img_white = Image.open(white_bg_path).convert('RGB')

    # Resize white to match black if dimensions differ
    if img_black.size != img_white.size:
        print(f"Warning: Size mismatch. Resizing white ({img_white.size}) to match black ({img_black.size})")
        img_white = img_white.resize(img_black.size, Image.LANCZOS)

    width, height = img_black.size

    # Convert to numpy arrays for faster processing
    data_black = np.array(img_black, dtype=np.float64)
    data_white = np.array(img_white, dtype=np.float64)

    # Distance between white (255,255,255) and black (0,0,0)
    # sqrt(255^2 + 255^2 + 255^2) ≈ 441.67
    bg_dist = math.sqrt(3 * 255 * 255)

    # Calculate pixel distance between the two images
    diff = data_white - data_black
    pixel_dist = np.sqrt(np.sum(diff ** 2, axis=2))

    # Calculate alpha: opaque pixels have distance 0, transparent pixels have max distance
    alpha = 1.0 - (pixel_dist / bg_dist)
    alpha = np.clip(alpha, 0, 1)

    # Color recovery: divide by alpha to un-premultiply
    # Use the black background image since C_observed = C_original * alpha + 0 * (1-alpha)
    # So C_original = C_observed / alpha
    alpha_safe = np.where(alpha > 0.01, alpha, 1.0)  # Avoid division by zero

    r_out = data_black[:, :, 0] / alpha_safe
    g_out = data_black[:, :, 1] / alpha_safe
    b_out = data_black[:, :, 2] / alpha_safe

    # Clamp to valid range
    r_out = np.clip(r_out, 0, 255)
    g_out = np.clip(g_out, 0, 255)
    b_out = np.clip(b_out, 0, 255)

    # Create output RGBA image
    output = np.zeros((height, width, 4), dtype=np.uint8)
    output[:, :, 0] = np.round(r_out).astype(np.uint8)
    output[:, :, 1] = np.round(g_out).astype(np.uint8)
    output[:, :, 2] = np.round(b_out).astype(np.uint8)
    output[:, :, 3] = np.round(alpha * 255).astype(np.uint8)

    # Save as PNG
    result = Image.fromarray(output, mode='RGBA')
    result.save(output_path, 'PNG')
    print(f"Saved transparent image to: {output_path}")


def main():
    if len(sys.argv) != 4:
        print("Usage: python stitch.py <black_bg_image> <white_bg_image> <output.png>")
        print("  black_bg_image: Image with the subject on a black background")
        print("  white_bg_image: Image with the subject on a white background")
        print("  output.png: Output file (will be a transparent PNG)")
        sys.exit(1)

    black_bg_path = sys.argv[1]
    white_bg_path = sys.argv[2]
    output_path = sys.argv[3]

    extract_alpha(black_bg_path, white_bg_path, output_path)


if __name__ == '__main__':
    main()