-
Notifications
You must be signed in to change notification settings - Fork 3
/
preprocess.py
161 lines (123 loc) · 4.21 KB
/
preprocess.py
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
153
154
155
156
157
158
159
160
161
# coding=utf-8
"""
@Author: LiYangMing
@StartTime: 18/12/29
@Filename: preprocess
@Software: Pycharm
@LastModify: 18/12/31
"""
from pyspark import SparkContext
from pyspark.streaming import StreamingContext
import re
import os
import tqdm
import time
from PIL import Image
class ImageProcessor(object):
"""
基于 PIL 库的图像预处理.
"""
@staticmethod
def resize_file(img, width, height):
"""
重新调整一个 Image 的大小.
"""
return img.resize((width, height), Image.BILINEAR)
@staticmethod
def single_channel(img):
return img.convert("L")
@staticmethod
def convert_dir(in_dir, out_dir):
"""
批量的处理一个目录下的所有图像.
"""
if not os.path.exists(out_dir):
os.system("mkdir -p " + out_dir)
for img_file in tqdm.tqdm(os.listdir(in_dir)):
img_path = os.path.join(in_dir, img_file)
r_image = Image.open(img_path)
c_image = ImageProcessor.resize_file(
r_image, 32, 32
).convert("RGBA")
out_path = os.path.join(out_dir, img_file)
c_image.save(out_path.split(".")[0] + ".png")
class FileStream(object):
"""
流批量 Spark Streaming 处理文件对象.
"""
def __init__(self, sc):
self._sc = sc
@staticmethod
def get_all_path(dir_list):
path_list = []
for directory in dir_list:
for file_name in os.listdir(directory):
file_path = os.path.join(directory, file_name)
path_list.append(file_path)
return path_list
@staticmethod
def bucket_segment(path_list, n_bucket):
total_num = len(path_list)
bucket_len = total_num // (n_bucket + 1) + 1
segment_list = []
for start in range(0, total_num, bucket_len):
end = min(start + bucket_len, total_num)
segment_list.append(path_list[start: end])
return segment_list
def stream_group(self, dir_list, out_dir, n_bucket):
if not os.path.exists(out_dir):
os.system("mkdir -p " + out_dir)
# 对文件路径做切分.
path_list = FileStream.get_all_path(dir_list)
seg_path_list = FileStream.bucket_segment(path_list, n_bucket)
# 初始化一个流式处理对象.
ssc = StreamingContext(self._sc, 1)
# 将路径列表压入队列.
queue_rdd = ssc.queueStream(seg_path_list, oneAtATime=True)
# RDD 的转换过程如下, 要综合考察各种可能的情况.
# 例如 3.png, 12.jpg, 7f.jpg, 9-over.png.
pair_rdd = queue_rdd.map(
lambda path: (re.findall("\d+", os.path.split(path)[-1])[0], path)
)
group_rdd = pair_rdd.groupByKey()
def _dump_image(x):
(index, p_list) = x
w_dir = os.path.join(out_dir, str(index))
if not os.path.exists(w_dir):
os.system("mkdir -p " + w_dir)
out_list = []
for r_path in p_list:
name = ("-".join(r_path.split("-")[1:])).replace("/", "$")
w_path = os.path.join(w_dir, name)
image = Image.open(r_path)
image.save(w_path)
out_list.append(w_path)
return out_list
group_rdd.map(_dump_image).pprint()
ssc.start()
time.sleep(40)
ssc.stop(stopSparkContext=False)
if __name__ == '__main__':
# 修改图片的格式, 包括尺寸, 通道.
# ImageProcessor.convert_dir(
# "data/material/pokemon-a",
# "data/pokemon/pokemon-a"
# )
# ImageProcessor.convert_dir(
# "data/material/pokemon-b",
# "data/pokemon/pokemon-b"
# )
# ImageProcessor.convert_dir(
# "data/material/pokemon-tcg-images",
# "data/pokemon/pokemon-tcg-images"
# )
sc = SparkContext("local[1]", "Group Image")
file_stream = FileStream(sc)
# file_stream.stream_group(
# ["data/pokemon/pokemon-a", "data/pokemon/pokemon-b"],
# "data/train", n_bucket=16
# )
file_stream.stream_group(
["data/pokemon/pokemon-tcg-images"],
"data/test", n_bucket=16
)