Search code examples
pythonalgorithmpython-imaging-librarychromakey

How to trim out a green screen and crop an image as fast as possible?


So I've been trying to remove the green screen, and crop, and I've had success, however, it is pretty slow, especially when I'm trying to use it on hundreds of pictures. I am not very familiar with image processing or PIL libraries, so I would love advice on how to make my code faster.

How it works: Basically it loops through each pixel, recording the pixel it looped over, and does it until it hits a non green like color, at which point it records the number of pixels away from the edge. I went with four loops, because i wanted to minimize the number of pixels i had to traverse (I can do the same thing with one loop but it would traverse across every pixel). The visitedPixel set prevents dealing with the same pixel. After the loops were done, it got a set of pixels that can be used to trim out the green screen edges, and thus cropping the image.

def trim_greenscreen_and_crop(image_name, output_name):

    img = Image.open(image_name)
    pixels = img.load()
    width = img.size[0]
    height = img.size[1]
    visitedPixel = set()
    box = [0, 0, 0, 0]

    # left edge
    break_flag = False
    for x in range(width):
        for y in range(height):
            coords = (x, y)
            r, g, b = pixels[x, y]
            if not (g > r and g > b and g > 200) and coords not in visitedPixel:
                box[0] = x - 1
                break_flag = True
                break
            visitedPixel.add(coords)
        if break_flag:
            break

    # top edge
    break_flag = False
    for y in range(height):
        for x in range(width):
            coords = (x, y)
            r, g, b = pixels[x, y]
            if not (g > r and g > b and g > 200) and coords not in visitedPixel:
                box[1] = y-1
                break_flag = True
                break
            visitedPixel.add(coords)
        if break_flag:
            break

    # right edge
    break_flag = False
    for x in range(width - 1, -1, -1):
        for y in range(height):
            coords = (x, y)
            r, g, b = pixels[x, y]
            if not (g > r and g > b and g > 200) and coords not in visitedPixel:
                box[2] = x + 1
                break_flag = True
                break
            visitedPixel.add(coords)
        if break_flag:
            break

    # bottom edge
    break_flag = False
    for y in range(height - 1, -1, -1):
        for x in range(width):
            coords = (x, y)
            r, g, b = pixels[x, y]
            if not (g > r and g > b and g > 200) and coords not in visitedPixel:
                box[3] = y + 1
                break_flag = True
                break
            visitedPixel.add(coords)
        if break_flag:
            break
    cropped_img = img.crop(box)

    if cropped_img.size == (0, 0):
        return img.size

    # cropped_img.save(output_name)
    return cropped_img.size

Before:

Before

After:

enter image description here


Solution

  • So i figured using numpy, and got this much faster solution which involves finding the variance of the rows and columns, thanks to MarkSetchell's idea.

    draft:

    def trim_greenscreen_and_crop(image_name, output_name):
        # use numpy to read the image
        img = Image.open(image_name)
        np_img = np.array(Image.open(image_name))
        # use numpy to get the variance across the rows and columns
        row_var = np.var(np_img, axis=0)
        col_var = np.var(np_img, axis=1)
    
        # select the rows and columns with some variance (basically not all green)
        no_variance_row = np.where(row_var > 5)
        no_variance_col = np.where(col_var > 5)
    
        # checks if the entire image is green, then dont trim
        if len(no_variance_row[0]) == 0 or len(no_variance_col[0]) == 0:
            return img.size
        else:
            # crops the image using the distance from the edges to the first non-green pixel
            cropped_img = img.crop((no_variance_row[0][0], no_variance_col[0][0], no_variance_row[0][-1], no_variance_col[0][-1]))
            cropped_img.save(output_name)
        return cropped_img.size