Simple image data utilities.

crop_imgs.py

#!/usr/bin/python
# -*- coding: utf-8 -*-
import glob
import os
import time
import smartcrop
from PIL import Image
from multiprocessing import Pool
import pickle


INPUT_GLOB = 'raw/*/*'
EXIST = set(glob.glob('data/*'))
N_PROCS = 8


def proc(fn):
    base_fn = fn.replace('\\', '_').replace('/', '_') + '.png'
    save_fn = os.path.join('data', base_fn)
    if save_fn in EXIST or os.path.exists(save_fn):
        return
    try:
        img = Image.open(fn)
        img = img.convert('RGB')
        sc = smartcrop.SmartCrop()
        w, h = img.size
        result = sc.crop(img, min(w, h) - 20, min(w, h) - 20)
        crop = result['top_crop']
        x = crop['x']
        y = crop['y']
        w = crop['width']
        h = crop['height']
        img = img.crop((x, y, x + w, y + h))
        img = img.resize((256, 256))
        img.save(save_fn)
    except Exception as e:
        print(fn, e)


if __name__ == '__main__':
    with Pool(N_PROCS) as pool:
        pool.map(proc, glob.glob(INPUT_GLOB))

label_imgs.py

#!/usr/bin/python
# -*- coding: utf-8 -*-
import glob
import os
import cv2
import random


GLOB_PATHS = [r'data/*.png']
SIZE = (256, 256)
KEYS = {ord(x): x for x in 'qwertyuiopasdfghjklzxcvbnm'}
LABELS_FN = 'labels.flist'


def main():

    if os.path.exists(LABELS_FN):
        with open(LABELS_FN, 'r') as f:
            data = f.read()
            history = set(item.split(',')[1] for item in
                          data.strip().split('\n') if len(item) > 0)
    else:
        history = set()

    fns = []
    for path in GLOB_PATHS:
        fns.extend(glob.glob(path))
    random.shuffle(fns)

    for fn in fns:
        if fn in history:
            continue
        img = cv2.resize(cv2.imread(fn), (256, 256))
        while True:
            cv2.imshow('img', img)
            key = cv2.waitKey(1) & 0xff
            if key in KEYS:
                label = KEYS[key]
                with open(LABELS_FN, 'a') as f:
                    f.write(label + ',' + fn + '\n')
                history.add(fn)
                break
            elif key == ord(' '):
                done = True
                break


if __name__ == '__main__':
    main()