def crop_face(args):
assert osp.exists(args.data_dir), "The input dir not exist"
root_folder_name = args.data_dir.split('/')[-1]
dst_folder = args.data_dir.replace(root_folder_name, root_folder_name + '_OPPOFaces')
lz.mkdir_p(dst_folder, delete=False)
mtcnn = MTCNN()
ind = 0
all_img = []
for imgfn in itertools.chain(
glob.glob(args.data_dir + '/**/*.jpg', recursive=True),
glob.glob(args.data_dir + '/**/*.JPEG', recursive=True)):
ind += 1
if ind % 10 == 0:
print(f'proc {ind}, {imgfn}')
dstimgfn = imgfn.replace(root_folder_name, root_folder_name + '_OPPOFaces')
dst_folder = osp.dirname(dstimgfn)
lz.mkdir_p(dst_folder, delete=False)
img = cvb.read_img(imgfn) # bgr
img1 = Image.fromarray(img)
face = mtcnn.align_best(img1, limit=None, min_face_size=16, imgfn=imgfn)
face = np.asarray(face) # bgr
# face = cvb.bgr2rgb(face) # rgb
cvb.write_img(face, dstimgfn)
all_img.append( dstimgfn)
logging.info(f'finish crop all {ind} imgs')
lz.msgpack_dump(all_img, dst_folder + '/' + 'all_imgs.pk')
del mtcnn
torch.cuda.empty_cache()
def consumer(queue, lock):
while True:
imgfn, param, roi_box, dst_imgfn = queue.get()
pts68 = [predict_68pts(param[i], roi_box[i]) for i in range(param.shape[0])]
for img_fp, pts68_, dst in zip(imgfn, pts68, dst_imgfn):
try:
img_ori = cvb.read_img(img_fp)
pts5 = to_landmark5(pts68_[:2, :].transpose())
warped = preprocess(img_ori, landmark=pts5)
# plt_imshow(warped, inp_mode='bgr'); plt.show()
lz.mkdir_p(osp.dirname(dst), delete=False)
cvb.write_img(warped, dst)
except Exception as e:
logging.warning(f'error occur {e}, pls check!')
cvb.write_img(np.ones((112, 112, 3), np.uint8), dst)
def test_write_img(self):
img = cvb.read_img(self.img_path)
out_file = '.cvbase_test.tmp.jpg'
cvb.write_img(img, out_file)
assert cvb.read_img(out_file).shape == (300, 400, 3)
os.remove(out_file)