bg.py

import os
import sys
import matplotlib.pyplot as plt
import matplotlib.image
import scipy.ndimage as ndimage
import numpy as np

width, height = 1280, 720
grid = 20
# Difference (norm2) between old and new pixels to consider it's not the background
pixel_threshold = 0.08
# Number of different pixel to consider it's not the background
grid_threshold = grid * grid * 0.3
# Ratio of new image incorporated in the old background
improvement_rate = 0.2

learn_dir = "./learn"
test_dir = "./test"

# === Learning process: compute the background ===

background = np.zeros((height, width, 3), dtype=float)
i = 0
for filename in os.listdir(learn_dir):
    if filename.endswith(".jpg"):
        i += 1
        background += matplotlib.image.imread(os.path.join(learn_dir, filename)) / 255

if i == 0:
    raise Exception("Empty learn directory")

background /= i

# plt.figure()
# plt.imshow(background)
# plt.show()

# === Test process: extract the foreground ===

for filename in os.listdir(test_dir):
    if not filename.endswith(".jpg"):
        continue
    image = matplotlib.image.imread(os.path.join(test_dir, filename)) / 255
    alpha = np.zeros((height, width), dtype=float)
    cells_to_keep = np.zeros((int(height / grid), int(width / grid)), dtype=bool)
    same_pixels = np.linalg.norm(image - background, axis=2) < pixel_threshold
    for y in range(int(height / grid)):
        for x in range(int(width / grid)):
            keep = (
                np.sum(
                    same_pixels[y * grid : y * grid + grid, x * grid : x * grid + grid],
                    axis=(0, 1),
                )
                <= grid_threshold
            )
            cells_to_keep[y, x] = keep
    # Add inner cells that were removed, and remove lone cells
    for y in range(int(height / grid)):
        for x in range(int(width / grid)):
            keep = cells_to_keep[y, x]
            if (
                not keep
                and np.sum(cells_to_keep[y - 1 : y + 2, x - 1 : x + 2], axis=(0, 1))
                >= 5
            ):
                keep = True
            elif (
                keep
                and np.sum(cells_to_keep[y - 1 : y + 2, x - 1 : x + 2], axis=(0, 1))
                <= 1
            ):
                keep = False
            cells_to_keep[y, x] = keep
            alpha[y * grid : y * grid + grid, x * grid : x * grid + grid] = float(
                cells_to_keep[y, x]
            )
            # Improve the background
            if not keep:
                background[y * grid : y * grid + grid, x * grid : x * grid + grid] *= (
                    1 - improvement_rate
                )
                background[y * grid : y * grid + grid, x * grid : x * grid + grid] += (
                    improvement_rate
                    * image[y * grid : y * grid + grid, x * grid : x * grid + grid]
                )

    plt.figure()
    plt.imshow(image)
    # Gaussian blur to make sharp edges smooth
    alpha = ndimage.gaussian_filter(alpha, sigma=grid)
    plt.imshow(np.zeros((height, width), dtype=float), alpha=(1 - alpha))
    plt.show()