darkliang
102 lines
4.1 KiB
Python
102 lines
4.1 KiB
Python
import os
|
|
import pickle
|
|
import numpy as np
|
|
import cv2
|
|
from tqdm import tqdm
|
|
import argparse
|
|
|
|
|
|
T_W = 64
|
|
T_H = 64
|
|
|
|
|
|
def cut_img(img):
|
|
# A silhouette contains too little white pixels
|
|
# might be not valid for identification.
|
|
# Get the top and bottom point
|
|
y = img.sum(axis=1)
|
|
y_top = (y != 0).argmax(axis=0)
|
|
y_btm = (y != 0).cumsum(axis=0).argmax(axis=0)
|
|
img = img[y_top:y_btm + 1, :]
|
|
# As the height of a person is larger than the width,
|
|
# use the height to calculate resize ratio.
|
|
_r = img.shape[1] / img.shape[0]
|
|
_t_w = int(T_H * _r)
|
|
img = cv2.resize(img, (_t_w, T_H), interpolation=cv2.INTER_AREA)
|
|
# Get the median of x axis and regard it as the x center of the person.
|
|
sum_point = img.sum()
|
|
sum_column = img.sum(axis=0).cumsum()
|
|
x_center = -1
|
|
for i in range(sum_column.size):
|
|
if sum_column[i] > sum_point / 2:
|
|
x_center = i
|
|
break
|
|
if x_center < 0:
|
|
return None
|
|
h_T_W = int(T_W / 2)
|
|
left = x_center - h_T_W
|
|
right = x_center + h_T_W
|
|
if left <= 0 or right >= img.shape[1]:
|
|
left += h_T_W
|
|
right += h_T_W
|
|
_ = np.zeros((img.shape[0], h_T_W))
|
|
img = np.concatenate([_, img, _], axis=1)
|
|
img = img[:, left:right]
|
|
return img.astype('uint8')
|
|
|
|
|
|
if __name__ == '__main__':
|
|
parser = argparse.ArgumentParser(description='CCPG dataset Preprocessing.')
|
|
parser.add_argument('--sil_path', default='', type=str,
|
|
help='Root path of raw silhouette dataset.')
|
|
parser.add_argument('--rgb_path', default='', type=str,
|
|
help='Root path of raw RGB dataset.')
|
|
parser.add_argument('-o', '--output_path', default='',
|
|
type=str, help='Output path of pickled dataset.')
|
|
args = parser.parse_args()
|
|
|
|
RGB_SIZE = (128, 128)
|
|
for _id in tqdm(sorted(os.listdir(args.sil_path))):
|
|
for _type in sorted(os.listdir(os.path.join(args.rgb_path, _id))):
|
|
for _view in sorted(os.listdir(os.path.join(args.rgb_path, _id, _type))):
|
|
imgs = []
|
|
segs = []
|
|
ratios = []
|
|
aligned_segs = []
|
|
for img_file in sorted(os.listdir(os.path.join(args.rgb_path, _id, _type, _view))):
|
|
seg_file = img_file.split(".")[0]+".png"
|
|
img_path = os.path.join(
|
|
args.rgb_path, _id, _type, _view, img_file)
|
|
seg_path = os.path.join(
|
|
args.rgb_path, _id, _type, _view, seg_file)
|
|
if not os.path.exists(seg_path):
|
|
print("Not Found: "+seg_path)
|
|
continue
|
|
img = cv2.imread(img_path)
|
|
seg = cv2.imread(seg_path, cv2.IMREAD_GRAYSCALE)
|
|
ratio = img.shape[1]/img.shape[0]
|
|
aligned_seg = cut_img(seg)
|
|
img = np.transpose(cv2.cvtColor(cv2.resize(
|
|
img, RGB_SIZE), cv2.COLOR_BGR2RGB), (2, 0, 1))
|
|
imgs.append(img)
|
|
segs.append(cv2.resize(
|
|
seg, RGB_SIZE))
|
|
aligned_segs.append(aligned_seg)
|
|
ratios.append(ratio)
|
|
if len(imgs) > 0:
|
|
output_path = os.path.join(
|
|
args.output_path, _id, _type, _view)
|
|
os.makedirs(output_path, exist_ok=True)
|
|
pickle.dump(np.asarray(imgs), open(os.path.join(
|
|
output_path, _view+"-rgbs.pkl"), "wb"))
|
|
pickle.dump(np.asarray(segs), open(os.path.join(
|
|
output_path, _view+"-sils.pkl"), "wb"))
|
|
pickle.dump(np.asarray(ratios), open(os.path.join(
|
|
output_path, _view+"-ratios.pkl"), "wb"))
|
|
pickle.dump(np.asarray(aligned_segs), open(os.path.join(
|
|
output_path, _view+"-aligned-sils.pkl"), "wb"))
|
|
else:
|
|
print("No imgs Found: " +
|
|
os.path.join(args.rgb_path, _id, _type, _view))
|
|
continue
|