utils/preprocess.py

import numpy as np
import os
import matplotlib.image as mpimage
import argparse
import functools
from utils import add_arguments, print_arguments
from dask.distributed import LocalCluster
from dask import bag as dbag
from dask.diagnostics import ProgressBar
from typing import Tuple
from PIL import Image
import cv2
#-----------------------------------#
#   对四个点坐标排序
#-----------------------------------#
def order_points(pts):
	# 一共4个坐标点
	rect = np.zeros((4, 2), dtype = "float32")

	# 按顺序找到对应坐标0123分别是 左上，右上，右下，左下
	# 计算左上，右下
	s = pts.sum(axis = 1)
	rect[0] = pts[np.argmin(s)]
	rect[2] = pts[np.argmax(s)]

	# 计算右上和左下
	diff = np.diff(pts, axis = 1)
	rect[1] = pts[np.argmin(diff)]
	rect[3] = pts[np.argmax(diff)]

	return rect
#-----------------------------------#
#   透射变换纠正车牌图片
#-----------------------------------#
def four_point_transform(image, pts):
	# 获取输入坐标点
	rect = order_points(pts)
	(tl, tr, br, bl) = rect

	# 计算输入的w和h值
	widthA = np.sqrt(((br[0] - bl[0]) ** 2) + ((br[1] - bl[1]) ** 2))
	widthB = np.sqrt(((tr[0] - tl[0]) ** 2) + ((tr[1] - tl[1]) ** 2))
	maxWidth = max(int(widthA), int(widthB))

	heightA = np.sqrt(((tr[0] - br[0]) ** 2) + ((tr[1] - br[1]) ** 2))
	heightB = np.sqrt(((tl[0] - bl[0]) ** 2) + ((tl[1] - bl[1]) ** 2))
	maxHeight = max(int(heightA), int(heightB))

	# 变换后对应坐标位置
	dst = np.array([
		[0, 0],
		[maxWidth - 1, 0],
		[maxWidth - 1, maxHeight - 1],
		[0, maxHeight - 1]], dtype = "float32")

	# 计算变换矩阵
	M = cv2.getPerspectiveTransform(rect, dst)
	warped = cv2.warpPerspective(image, M, (maxWidth, maxHeight))

	# 返回变换后结果
	return warped

# Dataset statistics that I gathered in development
#-----------------------------------#
#   用于过滤感知质量较低的不良图片
#-----------------------------------#
IMAGE_MEAN = 0.5
IMAGE_MEAN_STD = 0.028

IMG_STD = 0.28
IMG_STD_STD = 0.01


def readImage(fileName: str) -> np.ndarray:
    image = mpimage.imread(fileName)
    return image


#-----------------------------------#
#   从文件名中提取车牌的坐标
#-----------------------------------#


def parseLabel(label: str) -> Tuple[np.ndarray, np.ndarray]:
    annotation = label.split('-')[3].split('_')
    coor1 = [int(i) for i in annotation[0].split('&')]
    coor2 = [int(i) for i in annotation[1].split('&')]
    coor3 = [int(i) for i in annotation[2].split('&')]
    coor4 = [int(i) for i in annotation[3].split('&')]
    coor = np.array([coor1, coor2, coor3, coor4])
    center = np.mean(coor, axis=0)
    return coor, center.astype(int)


#-----------------------------------#
#   根据车牌坐标裁剪出车牌图像
#-----------------------------------#
# def cropImage(image: np.ndarray, coor: np.ndarray, center: np.ndarray) -> np.ndarray:
#     image = four_point_transform(image, coor)
#     return image

def cropImage(image: np.ndarray, coor: np.ndarray, center: np.ndarray) -> np.ndarray:
    maxW = np.max(coor[:, 0] - center[0])  # max plate width
    maxH = np.max(coor[:, 1] - center[1])  # max plate height

    xWanted = [64, 128, 192, 256]
    yWanted = [32, 64, 96, 128]

    found = False
    for w, h in zip(xWanted, yWanted):
        if maxW < w//2 and maxH < h//2:
            maxH = h//2
            maxW = w//2
            found = True
            break
    if not found:  # plate too large, discard
        return np.array([])
    elif center[1]-maxH < 0 or center[1]+maxH >= image.shape[1] or \
            center[0]-maxW < 0 or center[0] + maxW >= image.shape[0]:
        return np.array([])
    else:
        return image[center[1]-maxH:center[1]+maxH, center[0]-maxW:center[0]+maxW]

#-----------------------------------#
#           保存车牌图片
#-----------------------------------#


def saveImage(image: np.ndarray, fileName: str, outDir: str) -> int:
    if image.shape[0] == 0:
        return 0
    else:
        imgShape = image.shape
        if imgShape[1] == 64:
            mpimage.imsave(os.path.join(outDir, '64_32', fileName), image)
        elif imgShape[1] == 128:
            mpimage.imsave(os.path.join(outDir, '128_64', fileName), image)
        elif imgShape[1] == 208:
            mpimage.imsave(os.path.join(outDir, '192_96', fileName), image)
        else: #resize large images
            image = Image.fromarray(image).resize((192, 96))
            image = np.asarray(image) # back to numpy array
            mpimage.imsave(os.path.join(outDir, '192_96', fileName), image)
        return 1


#-----------------------------------#
# 包装成一个函数，以便将处理区分到不同目录
#-----------------------------------#

def processImage(file: str, inputDir: str, outputDir: str, subFolder: str) -> int:
    result = parseLabel(file)
    filePath = os.path.join(inputDir,subFolder, file)
    image = readImage(filePath)
    plate = cropImage(image, result[0], result[1])
    if plate.shape[0] == 0:
        return 0
    mean = np.mean(plate/255.0)
    std = np.std(plate/255.0)
    # bad brightness
    if mean <= IMAGE_MEAN - 10*IMAGE_MEAN_STD or mean >= IMAGE_MEAN + 10*IMAGE_MEAN_STD:
        return 0
    # low contrast
    if std <= IMG_STD - 10*IMG_STD_STD:
        return 0
    status = saveImage(plate, file, outputDir)
    return status


def main(argv):
    jobNum = int(argv.jobNum)
    outputDir = argv.outputDir
    inputDir = argv.inputDir
    try:
        os.mkdir(outputDir)
        for shape in ['64_32', '128_64', '192_96']:
            os.mkdir(os.path.join(outputDir, shape))
    except OSError:
        pass  # path already exists
    client = LocalCluster(n_workers=jobNum, threads_per_worker=5)  # IO intensive, more threads
    print('* number of workers:{}, \n* input dir:{}, \n* output dir:{}\n\n'.format(jobNum, inputDir, outputDir))
    for subFolder in ['ccpd_green', 'ccpd_base', 'ccpd_db', 'ccpd_fn', 'ccpd_rotate', 'ccpd_tilt', 'ccpd_weather']:
        fileList = os.listdir(os.path.join(inputDir, subFolder))
        print('* {} images found in {}. Start processing ...'.format(len(fileList), subFolder))
        toDo = dbag.from_sequence(fileList, npartitions=jobNum*30).persist()  # persist the bag in memory
        toDo = toDo.map(processImage, inputDir, outputDir, subFolder)
        pbar = ProgressBar(minimum=2.0)
        pbar.register()  # register all computations for better tracking
        result = toDo.compute()
        print('* image cropped: {}. Done ...'.format(sum(result)))
    client.close()  # shut down the cluster


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description=__doc__)
    add_arg = functools.partial(add_arguments, argparser=parser)
    add_arg('jobNum',           int,    4,                                         '处理图片的线程数')
    add_arg('inputDir',         str,    'datasets/CCPD2020',                       '输入图片目录')
    add_arg('outputDir',        str,    'datasets/CCPD2020_new',                   '保存图片目录')
    args = parser.parse_args()
    print_arguments(args)
    main(args)