Extract thumbnails from a image of video thumbnails grid

thumbnail_extraction.py

#TODO: test each square  if it's actually a screen shot
import numpy as np
#from scipy.ndimage import gaussian_filter
from scipy.signal import convolve2d
from PIL import Image, ImageDraw
import matplotlib.pyplot as plt
import cv2

import logging
log = logging.getLogger(__name__)

def plot(title, x,y):
    plt.clf()
    plt.plot(x, y)
    plt.axhline(y=np.mean(y), color='r', linestyle='-')
    plt.title(title)
    plt.savefig(title + '.png')

def func(grad, off, size):
    k0 = grad[off::-size]
    k1 = grad[off+size::size]
    num = k0.shape[0] + k1.shape[0]
    if num > 0:
        return (np.sum(k0) + np.sum(k1))/(num**0.75)
    else:
        return 0

func_v = np.vectorize(func)
func_v.excluded.add(0)

def find_borders(grad, min_size = 150, max_gap = 15, sim = 0.75):
    total = grad.shape[0]    
    off_max = np.argmax(grad)
    if grad[off_max] < np.mean(grad) * 4:
        print('hello')
        return (-1, total + 1, 0)
    #print('off_max', off_max)

    range_size = np.arange(min_size, int(total/2)+1, dtype=int)
    fits_size = func_v(grad, off_max, range_size)
    if log.level == logging.DEBUG:
        plot('out/size', range_size, fits_size)    
    best_size_idx = np.argmax(fits_size)
    # if fits_size[best_size_idx] < np.mean(fits_size)*2:
    #     print(fits_size[best_size_idx], np.mean(fits_size)*2)
    #     return (-1, total + 1, 0)
    best_size = range_size[best_size_idx]

    if off_max - max_gap < 0:
        off_max += best_size
    if off_max + max_gap >= total:
        off_max -= best_size
    range_off1 = np.arange(off_max - max_gap, off_max + max_gap)
    fits_off1 = func_v(grad, range_off1, best_size)
    if log.level == logging.DEBUG:
        plot('out/off', range_off1, fits_off1)

    #print(fits_off1)
    best2_off1_idx = np.argpartition(fits_off1, -2)[-2:]
    #print(best2_off1_idx)
    #print(fits_off1)
    off0_idx = best2_off1_idx[0]
    off1_idx = best2_off1_idx[1]

    #filtered = best2_off1_idx[np.where(fits_off1[best2_off1_idx] > mean_off1 * 2)]
    if fits_off1[off0_idx] > fits_off1[off1_idx] * sim:
        # we have two
        best_off, best_gap =  (range_off1[max(off0_idx, off1_idx)], np.abs(off1_idx - off0_idx))
    else:
        best_off, best_gap = (range_off1[off1_idx], 0)

    while best_off - best_size >= 0:
         best_off -= best_size
    return (best_off, best_size, best_gap)


def extract(image, min_width=200, min_height=150, max_gap=15):
    width, height = image.size
    img = np.array(image)

    #img = gaussian_filter(img, 7)
    img = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
    #blur = cv2.GaussianBlur(img, (5, 5), 0.5)

    # grayx = np.linalg.norm(cv2.Scharr(img, cv2.CV_32F, 1, 0), axis=2)
    # grayy = np.linalg.norm(cv2.Scharr(img, cv2.CV_32F, 0, 1), axis=2)

    #grayx = np.linalg.norm(img[1:, :] - img[:-1, :])

    mygradx = np.array([[-3, -10, -3],
                        [3, 10, 3]])

    gradx = convolve2d(img, mygradx, mode='valid')
    grayx = gradx # np.linalg.norm(gradx, axis=2)
    grady = convolve2d(img, mygradx.transpose(), mode='valid')
    grayy = grady # np.linalg.norm(grady, axis=2)

    mean_grayx = np.mean(grayx)
    mean_grayy = np.mean(grayy)

    # cols = np.count_nonzero(grayx > mean_grayx * 2, axis=1)
    # rows = np.count_nonzero(grayy > mean_grayy * 2, axis=0)

    cols = np.sum(np.abs(grayx), axis=1)
    rows = np.sum(np.abs(grayy), axis=0)

    if log.level == logging.DEBUG:
        #Image.fromarray(grayx).show()
        #Image.fromarray(grayy).show()
        plot('out/cols', range(height-1), cols)
        plot('out/rows', range(width-1), rows)    

    mean_rows = np.mean(rows)
    mean_cols = np.mean(cols)
    idx_rows = np.argwhere(rows > mean_rows * 2)
    idx_cols = np.argwhere(cols > mean_cols * 2)


    off_x, size_x, gap_x = find_borders(rows, min_size=min_width, max_gap=max_gap)
    off_y, size_y, gap_y = find_borders(cols, min_size=min_height, max_gap=max_gap)

    log.debug(off_x, size_x, gap_x)
    log.debug(off_y, size_y, gap_y)
    if log.level == logging.DEBUG:
        draw = ImageDraw.Draw(image)
        off = off_x
        while off < width:
            draw.line([(off, 0), (off, height)], (0, 255, 0))
            off += size_x
            draw.line([((off - gap_x, 0)), (off - gap_x, height)], (0, 255, 0))
        
        off = off_y
        while off < height:
            draw.line([(0, off),(width, off)], (0, 255, 0))
            off += size_y
            draw.line([(0, off - gap_y), (width, off - gap_y)], (0, 255, 0))

        image.save('out/annote.png')

    x = 0
    new_x = off_x + 1
    while width - x >= size_x / 2:
        y = 0
        new_y = off_y + 1
        while height - y >= size_y / 2:
            if new_y - y - gap_y > size_y / 2 and new_x - x - gap_x > size_x / 2:
                log.debug(x, y, new_x, new_y, gap_x, gap_y)
                crop = image.crop((x, y, np.minimum(new_x - gap_x, width), np.minimum(new_y - gap_y, height)))
                yield crop
            y = new_y
            new_y += size_y
        x = new_x
        new_x += size_x
    
def test():
    log.setLevel(logging.DEBUG)
    image = Image.open('thumbnail2.jpg')
    for i, crop in enumerate(extract(image)):
        crop.save('out/extract'+ str(i) + '.jpg')

if __name__ == '__main__':
    test()

1. assumes that all borders are aligned.
2. detects padding between thumbnails as well.

An example of thumbnail grid can be found here, although the problem we solve here is the opposite.
https://unix.stackexchange.com/questions/63769/fast-tool-to-generate-thumbnail-video-galleries-for-command-line

Note that we used a slimmer Scharr operator - this is so because the 'edge' we want to detect between two thumbnails is NOT a pixel, but between two pixels. We can't cut a pixel in half!