def __call__(self, ant, stop_idx=-1):
# load image
img_name = self.get_img_name(ant)
img_path = self.get_img_path(self.imgs_root_path, img_name)
img = self.get_img(img_path)
# len(x_list)=98, len(y_list)=98
x_list, y_list = get_points_list(ant, stop_idx)
bod_map = get_bod_map(img,
x_list,
y_list,
resize_size=self.resize_size,
line_thickness=self.line_thickness,
gaussian_kernel=self.gaussian_kernel,
gaussian_sigma=self.gaussian_sigma)
bod_map_resized = get_bod_map(img,
x_list,
y_list,
resize_size=(256, 256),
line_thickness=self.line_thickness,
gaussian_kernel=self.gaussian_kernel,
gaussian_sigma=self.gaussian_sigma)
# For Visualization. by default (3, 64, 64)
bod_img = get_bod_img(img,
x_list,
y_list,
resize_size=self.resize_size,
line_thickness=self.line_thickness,
gaussian_kernel=self.gaussian_kernel,
gaussian_sigma=self.gaussian_sigma)
return img, bod_map, bod_img, img_name, bod_map_resized
def get_required_data(self, idx):
img_name = self.get_img_name(self.ant[idx])
img, bod_map, bod_map_resized = None, None, None
img_path = self.get_img_path(self.imgs_root_path, img_name)
_img = self.get_img(img_path) # (3, 256, 256)
if 'heatmap' in self._variables or 'resized_heatmap' in self._variables:
# len(x_list)=98, len(y_list)=98
x_list, y_list = get_points_list(self.ant[idx])
if 'image' in self._variables:
img = self.normalize_img(_img)
if 'heatmap' in self._variables:
bod_map = get_bod_map(_img, x_list, y_list, resize_size=self.resize_size, line_thickness=self.line_thickness,
gaussian_kernel=self.gaussian_kernel, gaussian_sigma=self.gaussian_sigma) # (15, 64, 64)
bod_map = self.normalize_img(bod_map, normalize_type="heatmap")
if 'resized_heatmap' in self._variables:
bod_map_resized = get_bod_map(_img, x_list, y_list, resize_size=(256, 256), line_thickness=self.line_thickness,
gaussian_kernel=self.gaussian_kernel, gaussian_sigma=self.gaussian_sigma) # (15, 256, 256)
bod_map_resized = self.normalize_img(
bod_map_resized, normalize_type="heatmap")
return [_ for _ in (img, bod_map, bod_map_resized) if _ is not None]
def preprocess_WFLW(args):
import csv
print("preprocessing WFLW dataset...")
src_dir = args.src_dir
assert os.path.isdir(src_dir)
out_dir = args.out_dir
os.makedirs(out_dir, exist_ok=True)
resize_size = args.resize_size
line_thickness = args.line_thickness
gaussian_kernel = args.gaussian_kernel
gaussian_sigma = args.gaussian_sigma
imgs_root_path = src_dir
assert os.path.exists(
imgs_root_path), f"specified path {imgs_root_path} not found."
out_csv = [["saved_name", "real_name"]]
mode = args.mode
textname = f"WFLW_annotations/list_98pt_rect_attr_train_test/list_98pt_rect_attr_{mode}.txt"
with open(os.path.join(src_dir, textname)) as f:
annotations = f.readlines()
annotations = [_.split(" ") for _ in annotations]
prep = Preprocessor(imgs_root_path, resize_size, line_thickness,
gaussian_kernel, gaussian_sigma)
tmp_hm_dict = dict()
tmp_img_dict = dict()
if args.save_boundary_image:
os.makedirs(os.path.join(out_dir, "WFLW_landmark_images", mode),
exist_ok=True)
os.makedirs(os.path.join(out_dir, "WFLW_cropped_images", mode),
exist_ok=True)
idx = 0
for annotation in tqdm(annotations):
img_name, img, y_list, x_list = get_croped_image(
annotation, os.path.join(src_dir, "WFLW_images"))
scale_ratio = 256. / img.shape[-1]
x_list_scaled = [int(_ * scale_ratio) for _ in x_list]
y_list_scaled = [int(_ * scale_ratio) for _ in y_list]
img_resized = imresize(img, (256, 256), channel_first=True)
bod_img = get_bod_img(img_resized, y_list_scaled, x_list_scaled,
resize_size, line_thickness, gaussian_kernel,
gaussian_sigma)
bod_map = get_bod_map(img_resized, y_list_scaled, x_list_scaled,
resize_size, line_thickness, gaussian_kernel,
gaussian_sigma)
saved_name = f"{mode}_{idx}.png"
tmp_img_dict[saved_name] = img_resized
tmp_hm_dict[saved_name] = bod_map # uint8
out_csv.append([saved_name, img_name])
if args.save_boundary_image:
save_path_bod = os.path.join(out_dir, "WFLW_landmark_images", mode,
saved_name)
save_path_cropped = os.path.join(out_dir, "WFLW_cropped_images",
mode, saved_name)
imsave(save_path_bod, bod_img, channel_first=True)
imsave(save_path_cropped, img_resized, channel_first=True)
idx += 1
np.savez_compressed(os.path.join(out_dir, f'WFLW_cropped_image_{mode}'),
**tmp_img_dict)
np.savez_compressed(os.path.join(out_dir, f'WFLW_heatmap_{mode}'),
**tmp_hm_dict)
with open(os.path.join(out_dir, f"{mode}_data.csv"), 'w') as f:
writer = csv.writer(f)
writer.writerows(out_csv)