iOSDC 2019 NOC heatmap 02

iosdc2019_heatmap_clients.py

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

import os
from PIL import Image, ImageDraw, ImageFont
import pyocr
import pyocr.builders


def read_clients_number(image_dir, num_list, left, top, width, height, ext='png'):
    clients_file = '{}/clients.csv'.format(image_dir)
    with open(clients_file, 'w') as cf:
        cf.write('file id, clients\n')
        for i in num_list:
            file_id = '{:04}'.format(i)
            org_file = '{}/org/{}.{}'.format(image_dir, file_id, ext)

            im = Image.open(org_file)
            im_clients_num = im.crop((left, top, left + width, top + height))
            im_clients_num = im_clients_num.convert('L')
            im_clients_num.point(lambda x: 1 if x != 0xFF else 0xFF)

            tools = pyocr.get_available_tools()
            txt = tools[0].image_to_string(
                im_clients_num,
                lang='eng',
                builder=pyocr.builders.TextBuilder(tesseract_layout=6))
            cf.write('{},{}\n'.format(file_id, txt))
            if txt is None or not txt.isdigit():
                print('failed (id:', i, ', result:', txt, ')')

def write_clients_number(image_dir, num_list, file_neme='clients', ext='png'):
    clients_file = '{}/{}.csv'.format(image_dir, file_neme)
    print(clients_file)
    with open(clients_file, 'r') as cf:
        cf.readline()
        for line in cf:
            data = line.strip().split(',')
            file_id, clients = data[0].strip(), data[1].strip()
            resize_file = '{}/resize/{}.{}'.format(image_dir, file_id, ext)
            with_clients_file = '{}/with_clients/{}.{}'.format(image_dir, file_id, ext)

            im_resize = Image.open(resize_file)
            draw = ImageDraw.Draw(im_resize)
            font = ImageFont.truetype("/usr/share/fonts/truetype/lato/Lato-Regular.ttf", 32)
            draw.text((5, 5), 'clients: {}'.format(clients), fill=(0, 0, 0), font=font)
            im_resize.save(with_clients_file)

if __name__ == '__main__':
    image_dir = '{}/data/{}'.format(os.environ['HOME'], 'iosdc2019_heatmap')
    read_clients_number(image_dir, range(1, 2795), 600, 25, 40, 30)
    write_clients_number(image_dir, range(1, 2795), file_neme='clients_fix')