mnist_util.py

import cv2 as cv
import numpy as np
import math


def _pad(img):
	""" Pad input image to get square image. """
	(h, w) = img.shape
	if w > h:
		diff = w - h
		return cv.copyMakeBorder(img, math.ceil(diff/2), math.floor(diff/2), 0, 0, cv.BORDER_CONSTANT)
	elif w < h:
		diff = h - w
		return cv.copyMakeBorder(img, 0 , 0, math.ceil(diff/2), math.floor(diff/2), cv.BORDER_CONSTANT)
	else:
		return img


def _downscale(img, size):
	""" Downscale input grayscale image with width or height (the largest) equal to "size". """
	(h, w) = img.shape
	if w > h:
		return cv.resize(img, (size, round((h * size)/w)))
	elif w < h:
		return cv.resize(img, (round((w*size)/h), size))
	else:
		return cv.resize(img, (size, size))


def _extract_clothing(img, contour):
	""" Extract clothing from background given corresponding contour. """
	im_extract = np.empty(img.shape, dtype=np.float32) 
	for i in range(img.shape[0]):
		for j in range(img.shape[1]):
			in_poly = cv.pointPolygonTest(contour, (j, i), False)
			im_extract[i,j] = img[i,j] if in_poly in [1, 0] else 0
	return im_extract


def _get_bounding_box(img):
	"""Get clothing contour and bounding rectangle. """
	threshold = 128
	ret, im_binary = cv.threshold(img, threshold, 255, cv.THRESH_BINARY_INV)

	contours, hierarchy = cv.findContours(im_binary, cv.RETR_TREE, cv.CHAIN_APPROX_SIMPLE)

	bound_rect = [cv.boundingRect(cv.approxPolyDP(c, 3, True)) for c in contours]
	rect_area = [r[2]*r[3] for r in bound_rect]
	
	largest_contour = contours[np.argmax(rect_area)]
	x, y, w, h = bound_rect[np.argmax(rect_area)]
	return largest_contour, (x, y, w, h)


def mnist_format(img):
	""" Convert input image to MNIST-like format. """
	im_gray = cv.cvtColor(img, cv.COLOR_BGR2GRAY)
	im_downscale = _downscale(im_gray, 300)

	contour, (x, y, w, h) = _get_bounding_box(im_downscale)

	im_inverse = cv.bitwise_not(im_downscale)
	im_extract = _extract_clothing(im_inverse, contour)

	im_crop = im_extract[y:y+h, x:x+w]
	
	im_square = _pad(im_crop)
	im_mnist = cv.resize(im_square, (28, 28))


	return im_mnist



def mnist_class_index(c):
	mnist_classes = {0: "T-shirt/top",
					 1: "Trouser",
					 2: "Pullover",
					 3: "Dress",
					 4: "Coat",
					 5: "Sandal",
					 6: "Shirt",
					 7: "Sneaker",
					 8: "Bag",
					 9: "Ankle boot"}
	return mnist_classes[c]