Hugging Face's logo

Spaces:
Egrt
/
LicenseGAN
Build error

App Files Community
LicenseGAN
File size: 5,381 Bytes 
905cd18
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
 
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



# 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:
    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:  # 车牌太大则丢弃
        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)
    # 亮度不好的
    if mean <= IMAGE_MEAN - 10*IMAGE_MEAN_STD or mean >= IMAGE_MEAN + 10*IMAGE_MEAN_STD:
        return 0
    # 低对比度的
    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  # 地址已经存在
    client = LocalCluster(n_workers=jobNum, threads_per_worker=5)  # 开启多线程
    for subFolder in ['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()  # 登记所有的计算，以便更好地跟踪
        result = toDo.compute()
        print('* image cropped: {}. Done ...'.format(sum(result)))
    client.close()  # 关闭集群


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/CCPD2019',                       '输入图片目录')
    add_arg('outputDir',        str,    'datasets/CCPD2019_new',                   '保存图片目录')
    args = parser.parse_args()
    print_arguments(args)
    main(args)