iOSDC 2019 NOC heatmap 01

iosdc2019_heatmap_crop.py

# -*- coding: utf-8 -*-

import os
from PIL import Image
import numpy as np


def find_edge(pix_line, pix_th):
    edge1, edge2 = 0, len(pix_line)
    for i, pix in enumerate(pix_line):
        if pix < pix_th:
            edge1 = i
            break
    for i, pix in enumerate(pix_line[::-1]):
        j = len(pix_line) - i
        if pix < pix_th:
            edge2 = j
            break
    return edge1, edge2


def crop_heatmap(image_dir, num_list, min_th, max_th, min_h, max_h, resize_width=720, ext='png'):
    for i in num_list:
        file_id = '{:04}'.format(i)
        org_file = '{}/org/{}.{}'.format(image_dir, file_id, ext)
        resize_file = '{}/resize/{}.{}'.format(image_dir, file_id, ext)

        im = Image.open(org_file)
        im_body = im.crop((0, 130, *im.size))
        im_w, im_h = im_body.size

        im_gray = im_body.convert('L')
        im_gray_np = np.array(im_gray)

        for pix_th in range(min_th, max_th, 0x08):
            left_edge, right_edge = find_edge(im_gray_np[int(im_h/2),:], pix_th)
            top_edge, bottom_edge = find_edge(im_gray_np[:,int(im_w/2)], pix_th)
            im_crop = im_body.crop((left_edge, top_edge, right_edge, bottom_edge))
            resize_height = round(im_crop.height * resize_width / im_crop.width)
            if min_h < resize_height and resize_height < max_h:
                break
        else:
            print('failed ( id:', i, 'size:', im_resize.height, ')')
        im_crop.resize((width, resize_height))
        im_crop.save(resize_file)


if __name__ == '__main__':
    image_dir = '{}/data/{}'.format(os.environ['HOME'], 'iosdc2019_heatmap')
    crop_heatmap(image_dir, range(1, 2795), 0x40, 0xC0, 515, 530)