Spaces:
Runtime error
Runtime error
import sys | |
import numpy as np | |
import cv2 | |
import os | |
from utils.tool import IoU,convert_to_square | |
import numpy.random as npr | |
import argparse | |
from utils.detect import MtcnnDetector, create_mtcnn_net | |
from utils.dataloader import ImageDB,TestImageLoader | |
import time | |
from six.moves import cPickle | |
import utils.config as config | |
import utils.vision as vision | |
sys.path.append(os.getcwd()) | |
txt_from_path = './data_set/wider_face_train_bbx_gt.txt' | |
anno_file = os.path.join(config.ANNO_STORE_DIR, 'anno_train.txt') | |
# anno_file = './anno_store/anno_train.txt' | |
prefix = '' | |
use_cuda = True | |
im_dir = "./data_set/face_detection/WIDER_train/images/" | |
traindata_store = './data_set/train/' | |
prefix_path = "./data_set/face_detection/WIDER_train/images/" | |
annotation_file = './anno_store/anno_train.txt' | |
prefix_path_lm = '' | |
annotation_file_lm = "./data_set/face_landmark/CNN_FacePoint/train/trainImageList.txt" | |
# ----------------------------------------------------other---------------------------------------------- | |
pos_save_dir = "./data_set/train/12/positive" | |
part_save_dir = "./data_set/train/12/part" | |
neg_save_dir = './data_set/train/12/negative' | |
pnet_postive_file = os.path.join(config.ANNO_STORE_DIR, 'pos_12.txt') | |
pnet_part_file = os.path.join(config.ANNO_STORE_DIR, 'part_12.txt') | |
pnet_neg_file = os.path.join(config.ANNO_STORE_DIR, 'neg_12.txt') | |
imglist_filename_pnet = os.path.join(config.ANNO_STORE_DIR, 'imglist_anno_12.txt') | |
# ----------------------------------------------------PNet---------------------------------------------- | |
rnet_postive_file = os.path.join(config.ANNO_STORE_DIR, 'pos_24.txt') | |
rnet_part_file = os.path.join(config.ANNO_STORE_DIR, 'part_24.txt') | |
rnet_neg_file = os.path.join(config.ANNO_STORE_DIR, 'neg_24.txt') | |
rnet_landmark_file = os.path.join(config.ANNO_STORE_DIR, 'landmark_24.txt') | |
imglist_filename_rnet = os.path.join(config.ANNO_STORE_DIR, 'imglist_anno_24.txt') | |
# ----------------------------------------------------RNet---------------------------------------------- | |
onet_postive_file = os.path.join(config.ANNO_STORE_DIR, 'pos_48.txt') | |
onet_part_file = os.path.join(config.ANNO_STORE_DIR, 'part_48.txt') | |
onet_neg_file = os.path.join(config.ANNO_STORE_DIR, 'neg_48.txt') | |
onet_landmark_file = os.path.join(config.ANNO_STORE_DIR, 'landmark_48.txt') | |
imglist_filename_onet = os.path.join(config.ANNO_STORE_DIR, 'imglist_anno_48.txt') | |
# ----------------------------------------------------ONet---------------------------------------------- | |
def assemble_data(output_file, anno_file_list=[]): | |
#assemble the pos, neg, part annotations to one file | |
size = 12 | |
if len(anno_file_list)==0: | |
return 0 | |
if os.path.exists(output_file): | |
os.remove(output_file) | |
for anno_file in anno_file_list: | |
with open(anno_file, 'r') as f: | |
print(anno_file) | |
anno_lines = f.readlines() | |
base_num = 250000 | |
if len(anno_lines) > base_num * 3: | |
idx_keep = npr.choice(len(anno_lines), size=base_num * 3, replace=True) | |
elif len(anno_lines) > 100000: | |
idx_keep = npr.choice(len(anno_lines), size=len(anno_lines), replace=True) | |
else: | |
idx_keep = np.arange(len(anno_lines)) | |
np.random.shuffle(idx_keep) | |
chose_count = 0 | |
with open(output_file, 'a+') as f: | |
for idx in idx_keep: | |
# write lables of pos, neg, part images | |
f.write(anno_lines[idx]) | |
chose_count+=1 | |
return chose_count | |
def wider_face(txt_from_path, txt_to_path): | |
line_from_count = 0 | |
with open(txt_from_path, 'r') as f: | |
annotations = f.readlines() | |
with open(txt_to_path, 'w+') as f: | |
while line_from_count < len(annotations): | |
if annotations[line_from_count][2]=='-': | |
img_name = annotations[line_from_count][:-1] | |
line_from_count += 1 # change line to read the number | |
bbox_count = int(annotations[line_from_count]) # num of bboxes | |
line_from_count += 1 # change line to read the posession | |
for _ in range(bbox_count): | |
bbox = list(map(int,annotations[line_from_count].split()[:4])) # give a loop to append all the boxes | |
bbox = [bbox[0], bbox[1], bbox[0]+bbox[2], bbox[1]+bbox[3]] # make x1, y1, w, h --> x1, y1, x2, y2 | |
bbox = list(map(str,bbox)) | |
img_name += (' '+' '.join(bbox)) | |
line_from_count+=1 | |
f.write(img_name +'\n') | |
else: # dectect the file name | |
line_from_count+=1 | |
# ----------------------------------------------------origin---------------------------------------------- | |
def get_Pnet_data(): | |
if not os.path.exists(pos_save_dir): | |
os.makedirs(pos_save_dir) | |
if not os.path.exists(part_save_dir): | |
os.makedirs(part_save_dir) | |
if not os.path.exists(neg_save_dir): | |
os.makedirs(neg_save_dir) | |
f1 = open(os.path.join('./anno_store', 'pos_12.txt'), 'w') | |
f2 = open(os.path.join('./anno_store', 'neg_12.txt'), 'w') | |
f3 = open(os.path.join('./anno_store', 'part_12.txt'), 'w') | |
with open(anno_file, 'r') as f: | |
annotations = f.readlines() | |
num = len(annotations) | |
print("%d pics in total" % num) | |
p_idx = 0 # positive | |
n_idx = 0 # negative | |
d_idx = 0 # dont care | |
idx = 0 | |
box_idx = 0 | |
for annotation in annotations: | |
annotation = annotation.strip().split(' ') | |
# annotation[0]文件名 | |
im_path = os.path.join(im_dir, annotation[0]) | |
# print(im_path) | |
# print(os.path.exists(im_path)) | |
bbox = list(map(float, annotation[1:])) | |
# annotation[1:]人脸坐标,一张脸4个值,对应两个点的坐标 | |
boxes = np.array(bbox, dtype=np.int32).reshape(-1, 4) | |
# -1处的值为人脸数目 | |
if boxes.shape[0]==0: | |
continue | |
# 若无人脸则跳过本次循环 | |
img = cv2.imread(im_path) | |
# print(img.shape) | |
# exit() | |
# 计数 | |
idx += 1 | |
if idx % 100 == 0: | |
print("%s images done, pos: %s part: %s neg: %s" % (idx, p_idx, d_idx, n_idx)) | |
# 图片三通道 | |
height, width, channel = img.shape | |
neg_num = 0 | |
# 取50次不同的框 | |
while neg_num < 50: | |
size = np.random.randint(12, min(width, height) / 2) | |
nx = np.random.randint(0, width - size) | |
ny = np.random.randint(0, height - size) | |
crop_box = np.array([nx, ny, nx + size, ny + size]) | |
Iou = IoU(crop_box, boxes) # IoU为 重合部分 / 两框之和 ,越大越好 | |
cropped_im = img[ny: ny + size, nx: nx + size, :] # 裁去多余部分并resize成 12*12 | |
resized_im = cv2.resize(cropped_im, (12, 12), interpolation=cv2.INTER_LINEAR) | |
if np.max(Iou) < 0.3: | |
# Iou with all gts must below 0.3 | |
save_file = os.path.join(neg_save_dir, "%s.jpg" % n_idx) | |
f2.write(save_file + ' 0\n') | |
cv2.imwrite(save_file, resized_im) | |
n_idx += 1 | |
neg_num += 1 | |
for box in boxes: | |
# box (x_left, y_top, x_right, y_bottom) | |
x1, y1, x2, y2 = box | |
# w = x2 - x1 + 1 | |
# h = y2 - y1 + 1 | |
w = x2 - x1 + 1 | |
h = y2 - y1 + 1 | |
# ignore small faces | |
# in case the ground truth boxes of small faces are not accurate | |
if max(w, h) < 40 or x1 < 0 or y1 < 0: | |
continue | |
if w < 12 or h < 12: | |
continue | |
# generate negative examples that have overlap with gt | |
for i in range(5): | |
size = np.random.randint(12, min(width, height) / 2) | |
# delta_x and delta_y are offsets of (x1, y1) | |
delta_x = np.random.randint(max(-size, -x1), w) | |
delta_y = np.random.randint(max(-size, -y1), h) | |
nx1 = max(0, x1 + delta_x) | |
ny1 = max(0, y1 + delta_y) | |
if nx1 + size > width or ny1 + size > height: | |
continue | |
crop_box = np.array([nx1, ny1, nx1 + size, ny1 + size]) | |
Iou = IoU(crop_box, boxes) | |
cropped_im = img[ny1: ny1 + size, nx1: nx1 + size, :] | |
resized_im = cv2.resize(cropped_im, (12, 12), interpolation=cv2.INTER_LINEAR) | |
if np.max(Iou) < 0.3: | |
# Iou with all gts must below 0.3 | |
save_file = os.path.join(neg_save_dir, "%s.jpg" % n_idx) | |
f2.write(save_file + ' 0\n') | |
cv2.imwrite(save_file, resized_im) | |
n_idx += 1 | |
# generate positive examples and part faces | |
for i in range(20): | |
size = np.random.randint(int(min(w, h) * 0.8), np.ceil(1.25 * max(w, h))) | |
# delta here is the offset of box center | |
delta_x = np.random.randint(-w * 0.2, w * 0.2) | |
delta_y = np.random.randint(-h * 0.2, h * 0.2) | |
nx1 = max(x1 + w / 2 + delta_x - size / 2, 0) | |
ny1 = max(y1 + h / 2 + delta_y - size / 2, 0) | |
nx2 = nx1 + size | |
ny2 = ny1 + size | |
if nx2 > width or ny2 > height: | |
continue | |
crop_box = np.array([nx1, ny1, nx2, ny2]) | |
offset_x1 = (x1 - nx1) / float(size) | |
offset_y1 = (y1 - ny1) / float(size) | |
offset_x2 = (x2 - nx2) / float(size) | |
offset_y2 = (y2 - ny2) / float(size) | |
cropped_im = img[int(ny1): int(ny2), int(nx1): int(nx2), :] | |
resized_im = cv2.resize(cropped_im, (12, 12), interpolation=cv2.INTER_LINEAR) | |
box_ = box.reshape(1, -1) | |
if IoU(crop_box, box_) >= 0.65: | |
save_file = os.path.join(pos_save_dir, "%s.jpg" % p_idx) | |
f1.write(save_file + ' 1 %.2f %.2f %.2f %.2f\n' % (offset_x1, offset_y1, offset_x2, offset_y2)) | |
cv2.imwrite(save_file, resized_im) | |
p_idx += 1 | |
elif IoU(crop_box, box_) >= 0.4: | |
save_file = os.path.join(part_save_dir, "%s.jpg" % d_idx) | |
f3.write(save_file + ' -1 %.2f %.2f %.2f %.2f\n' % (offset_x1, offset_y1, offset_x2, offset_y2)) | |
cv2.imwrite(save_file, resized_im) | |
d_idx += 1 | |
box_idx += 1 | |
#print("%s images done, pos: %s part: %s neg: %s" % (idx, p_idx, d_idx, n_idx)) | |
f1.close() | |
f2.close() | |
f3.close() | |
def assembel_Pnet_data(): | |
anno_list = [] | |
anno_list.append(pnet_postive_file) | |
anno_list.append(pnet_part_file) | |
anno_list.append(pnet_neg_file) | |
# anno_list.append(pnet_landmark_file) | |
chose_count = assemble_data(imglist_filename_pnet ,anno_list) | |
print("PNet train annotation result file path:%s" % imglist_filename_pnet) | |
# -----------------------------------------------------------------------------------------------------------------------------------------------# | |
def gen_rnet_data(data_dir, anno_file, pnet_model_file, prefix_path='', use_cuda=True, vis=False): | |
""" | |
:param data_dir: train data | |
:param anno_file: | |
:param pnet_model_file: | |
:param prefix_path: | |
:param use_cuda: | |
:param vis: | |
:return: | |
""" | |
# load trained pnet model | |
pnet, _, _ = create_mtcnn_net(p_model_path = pnet_model_file, use_cuda = use_cuda) | |
mtcnn_detector = MtcnnDetector(pnet = pnet, min_face_size = 12) | |
# load original_anno_file, length = 12880 | |
imagedb = ImageDB(anno_file, mode = "test", prefix_path = prefix_path) | |
imdb = imagedb.load_imdb() | |
image_reader = TestImageLoader(imdb, 1, False) | |
all_boxes = list() | |
batch_idx = 0 | |
print('size:%d' %image_reader.size) | |
for databatch in image_reader: | |
if batch_idx % 100 == 0: | |
print ("%d images done" % batch_idx) | |
im = databatch | |
t = time.time() | |
# obtain boxes and aligned boxes | |
boxes, boxes_align = mtcnn_detector.detect_pnet(im=im) | |
if boxes_align is None: | |
all_boxes.append(np.array([])) | |
batch_idx += 1 | |
continue | |
if vis: | |
rgb_im = cv2.cvtColor(np.asarray(im), cv2.COLOR_BGR2RGB) | |
vision.vis_two(rgb_im, boxes, boxes_align) | |
t1 = time.time() - t | |
print('cost time ',t1) | |
t = time.time() | |
all_boxes.append(boxes_align) | |
batch_idx += 1 | |
# if batch_idx == 100: | |
# break | |
# print("shape of all boxes {0}".format(all_boxes)) | |
# time.sleep(5) | |
# save_path = model_store_path() | |
# './model_store' | |
save_path = './model_store' | |
if not os.path.exists(save_path): | |
os.mkdir(save_path) | |
save_file = os.path.join(save_path, "detections_%d.pkl" % int(time.time())) | |
with open(save_file, 'wb') as f: | |
cPickle.dump(all_boxes, f, cPickle.HIGHEST_PROTOCOL) | |
# save_file = './model_store/detections_1588751332.pkl' | |
gen_rnet_sample_data(data_dir, anno_file, save_file, prefix_path) | |
def gen_rnet_sample_data(data_dir, anno_file, det_boxs_file, prefix_path): | |
""" | |
:param data_dir: | |
:param anno_file: original annotations file of wider face data | |
:param det_boxs_file: detection boxes file | |
:param prefix_path: | |
:return: | |
""" | |
neg_save_dir = os.path.join(data_dir, "24/negative") | |
pos_save_dir = os.path.join(data_dir, "24/positive") | |
part_save_dir = os.path.join(data_dir, "24/part") | |
for dir_path in [neg_save_dir, pos_save_dir, part_save_dir]: | |
# print(dir_path) | |
if not os.path.exists(dir_path): | |
os.makedirs(dir_path) | |
# load ground truth from annotation file | |
# format of each line: image/path [x1,y1,x2,y2] for each gt_box in this image | |
with open(anno_file, 'r') as f: | |
annotations = f.readlines() | |
image_size = 24 | |
net = "rnet" | |
im_idx_list = list() | |
gt_boxes_list = list() | |
num_of_images = len(annotations) | |
print ("processing %d images in total" % num_of_images) | |
for annotation in annotations: | |
annotation = annotation.strip().split(' ') | |
im_idx = os.path.join(prefix_path, annotation[0]) | |
# im_idx = annotation[0] | |
boxes = list(map(float, annotation[1:])) | |
boxes = np.array(boxes, dtype=np.float32).reshape(-1, 4) | |
im_idx_list.append(im_idx) | |
gt_boxes_list.append(boxes) | |
# './anno_store' | |
save_path = './anno_store' | |
if not os.path.exists(save_path): | |
os.makedirs(save_path) | |
f1 = open(os.path.join(save_path, 'pos_%d.txt' % image_size), 'w') | |
f2 = open(os.path.join(save_path, 'neg_%d.txt' % image_size), 'w') | |
f3 = open(os.path.join(save_path, 'part_%d.txt' % image_size), 'w') | |
# print(det_boxs_file) | |
det_handle = open(det_boxs_file, 'rb') | |
det_boxes = cPickle.load(det_handle) | |
# an image contain many boxes stored in an array | |
print(len(det_boxes), num_of_images) | |
# assert len(det_boxes) == num_of_images, "incorrect detections or ground truths" | |
# index of neg, pos and part face, used as their image names | |
n_idx = 0 | |
p_idx = 0 | |
d_idx = 0 | |
image_done = 0 | |
for im_idx, dets, gts in zip(im_idx_list, det_boxes, gt_boxes_list): | |
# if (im_idx+1) == 100: | |
# break | |
gts = np.array(gts, dtype=np.float32).reshape(-1, 4) | |
if gts.shape[0]==0: | |
continue | |
if image_done % 100 == 0: | |
print("%d images done" % image_done) | |
image_done += 1 | |
if dets.shape[0] == 0: | |
continue | |
img = cv2.imread(im_idx) | |
# change to square | |
dets = convert_to_square(dets) | |
dets[:, 0:4] = np.round(dets[:, 0:4]) | |
neg_num = 0 | |
for box in dets: | |
x_left, y_top, x_right, y_bottom, _ = box.astype(int) | |
width = x_right - x_left + 1 | |
height = y_bottom - y_top + 1 | |
# ignore box that is too small or beyond image border | |
if width < 20 or x_left < 0 or y_top < 0 or x_right > img.shape[1] - 1 or y_bottom > img.shape[0] - 1: | |
continue | |
# compute intersection over union(IoU) between current box and all gt boxes | |
Iou = IoU(box, gts) | |
cropped_im = img[y_top:y_bottom + 1, x_left:x_right + 1, :] | |
resized_im = cv2.resize(cropped_im, (image_size, image_size), | |
interpolation=cv2.INTER_LINEAR) | |
# save negative images and write label | |
# Iou with all gts must below 0.3 | |
if np.max(Iou) < 0.3 and neg_num < 60: | |
# save the examples | |
save_file = os.path.join(neg_save_dir, "%s.jpg" % n_idx) | |
# print(save_file) | |
f2.write(save_file + ' 0\n') | |
cv2.imwrite(save_file, resized_im) | |
n_idx += 1 | |
neg_num += 1 | |
else: | |
# find gt_box with the highest iou | |
idx = np.argmax(Iou) | |
assigned_gt = gts[idx] | |
x1, y1, x2, y2 = assigned_gt | |
# compute bbox reg label | |
offset_x1 = (x1 - x_left) / float(width) | |
offset_y1 = (y1 - y_top) / float(height) | |
offset_x2 = (x2 - x_right) / float(width) | |
offset_y2 = (y2 - y_bottom) / float(height) | |
# save positive and part-face images and write labels | |
if np.max(Iou) >= 0.65: | |
save_file = os.path.join(pos_save_dir, "%s.jpg" % p_idx) | |
f1.write(save_file + ' 1 %.2f %.2f %.2f %.2f\n' % ( | |
offset_x1, offset_y1, offset_x2, offset_y2)) | |
cv2.imwrite(save_file, resized_im) | |
p_idx += 1 | |
elif np.max(Iou) >= 0.4: | |
save_file = os.path.join(part_save_dir, "%s.jpg" % d_idx) | |
f3.write(save_file + ' -1 %.2f %.2f %.2f %.2f\n' % ( | |
offset_x1, offset_y1, offset_x2, offset_y2)) | |
cv2.imwrite(save_file, resized_im) | |
d_idx += 1 | |
f1.close() | |
f2.close() | |
f3.close() | |
def model_store_path(): | |
return os.path.dirname(os.path.dirname(os.path.dirname(os.path.realpath(__file__))))+"/model_store" | |
def get_Rnet_data(pnet_model): | |
gen_rnet_data(traindata_store, annotation_file, pnet_model_file = pnet_model, prefix_path = prefix_path, use_cuda = True) | |
def assembel_Rnet_data(): | |
anno_list = [] | |
anno_list.append(rnet_postive_file) | |
anno_list.append(rnet_part_file) | |
anno_list.append(rnet_neg_file) | |
# anno_list.append(pnet_landmark_file) | |
chose_count = assemble_data(imglist_filename_rnet ,anno_list) | |
print("RNet train annotation result file path:%s" % imglist_filename_rnet) | |
#-----------------------------------------------------------------------------------------------------------------------------------------------# | |
def gen_onet_data(data_dir, anno_file, pnet_model_file, rnet_model_file, prefix_path='', use_cuda=True, vis=False): | |
pnet, rnet, _ = create_mtcnn_net(p_model_path=pnet_model_file, r_model_path=rnet_model_file, use_cuda=use_cuda) | |
mtcnn_detector = MtcnnDetector(pnet=pnet, rnet=rnet, min_face_size=12) | |
imagedb = ImageDB(anno_file,mode="test",prefix_path=prefix_path) | |
imdb = imagedb.load_imdb() | |
image_reader = TestImageLoader(imdb,1,False) | |
all_boxes = list() | |
batch_idx = 0 | |
print('size:%d' % image_reader.size) | |
for databatch in image_reader: | |
if batch_idx % 50 == 0: | |
print("%d images done" % batch_idx) | |
im = databatch | |
t = time.time() | |
# pnet detection = [x1, y1, x2, y2, score, reg] | |
p_boxes, p_boxes_align = mtcnn_detector.detect_pnet(im=im) | |
t0 = time.time() - t | |
t = time.time() | |
# rnet detection | |
boxes, boxes_align = mtcnn_detector.detect_rnet(im=im, dets=p_boxes_align) | |
t1 = time.time() - t | |
print('cost time pnet--',t0,' rnet--',t1) | |
t = time.time() | |
if boxes_align is None: | |
all_boxes.append(np.array([])) | |
batch_idx += 1 | |
continue | |
if vis: | |
rgb_im = cv2.cvtColor(np.asarray(im), cv2.COLOR_BGR2RGB) | |
vision.vis_two(rgb_im, boxes, boxes_align) | |
all_boxes.append(boxes_align) | |
batch_idx += 1 | |
save_path = './model_store' | |
if not os.path.exists(save_path): | |
os.mkdir(save_path) | |
save_file = os.path.join(save_path, "detections_%d.pkl" % int(time.time())) | |
with open(save_file, 'wb') as f: | |
cPickle.dump(all_boxes, f, cPickle.HIGHEST_PROTOCOL) | |
gen_onet_sample_data(data_dir,anno_file,save_file,prefix_path) | |
def gen_onet_sample_data(data_dir,anno_file,det_boxs_file,prefix): | |
neg_save_dir = os.path.join(data_dir, "48/negative") | |
pos_save_dir = os.path.join(data_dir, "48/positive") | |
part_save_dir = os.path.join(data_dir, "48/part") | |
for dir_path in [neg_save_dir, pos_save_dir, part_save_dir]: | |
if not os.path.exists(dir_path): | |
os.makedirs(dir_path) | |
# load ground truth from annotation file | |
# format of each line: image/path [x1,y1,x2,y2] for each gt_box in this image | |
with open(anno_file, 'r') as f: | |
annotations = f.readlines() | |
image_size = 48 | |
net = "onet" | |
im_idx_list = list() | |
gt_boxes_list = list() | |
num_of_images = len(annotations) | |
print("processing %d images in total" % num_of_images) | |
for annotation in annotations: | |
annotation = annotation.strip().split(' ') | |
im_idx = os.path.join(prefix,annotation[0]) | |
boxes = list(map(float, annotation[1:])) | |
boxes = np.array(boxes, dtype=np.float32).reshape(-1, 4) | |
im_idx_list.append(im_idx) | |
gt_boxes_list.append(boxes) | |
save_path = './anno_store' | |
if not os.path.exists(save_path): | |
os.makedirs(save_path) | |
f1 = open(os.path.join(save_path, 'pos_%d.txt' % image_size), 'w') | |
f2 = open(os.path.join(save_path, 'neg_%d.txt' % image_size), 'w') | |
f3 = open(os.path.join(save_path, 'part_%d.txt' % image_size), 'w') | |
det_handle = open(det_boxs_file, 'rb') | |
det_boxes = cPickle.load(det_handle) | |
print(len(det_boxes), num_of_images) | |
# assert len(det_boxes) == num_of_images, "incorrect detections or ground truths" | |
# index of neg, pos and part face, used as their image names | |
n_idx = 0 | |
p_idx = 0 | |
d_idx = 0 | |
image_done = 0 | |
for im_idx, dets, gts in zip(im_idx_list, det_boxes, gt_boxes_list): | |
if image_done % 100 == 0: | |
print("%d images done" % image_done) | |
image_done += 1 | |
if gts.shape[0]==0: | |
continue | |
if dets.shape[0] == 0: | |
continue | |
img = cv2.imread(im_idx) | |
dets = convert_to_square(dets) | |
dets[:, 0:4] = np.round(dets[:, 0:4]) | |
for box in dets: | |
x_left, y_top, x_right, y_bottom = box[0:4].astype(int) | |
width = x_right - x_left + 1 | |
height = y_bottom - y_top + 1 | |
# ignore box that is too small or beyond image border | |
if width < 20 or x_left < 0 or y_top < 0 or x_right > img.shape[1] - 1 or y_bottom > img.shape[0] - 1: | |
continue | |
# compute intersection over union(IoU) between current box and all gt boxes | |
Iou = IoU(box, gts) | |
cropped_im = img[y_top:y_bottom + 1, x_left:x_right + 1, :] | |
resized_im = cv2.resize(cropped_im, (image_size, image_size), | |
interpolation=cv2.INTER_LINEAR) | |
# save negative images and write label | |
if np.max(Iou) < 0.3: | |
# Iou with all gts must below 0.3 | |
save_file = os.path.join(neg_save_dir, "%s.jpg" % n_idx) | |
f2.write(save_file + ' 0\n') | |
cv2.imwrite(save_file, resized_im) | |
n_idx += 1 | |
else: | |
# find gt_box with the highest iou | |
idx = np.argmax(Iou) | |
assigned_gt = gts[idx] | |
x1, y1, x2, y2 = assigned_gt | |
# compute bbox reg label | |
offset_x1 = (x1 - x_left) / float(width) | |
offset_y1 = (y1 - y_top) / float(height) | |
offset_x2 = (x2 - x_right) / float(width) | |
offset_y2 = (y2 - y_bottom) / float(height) | |
# save positive and part-face images and write labels | |
if np.max(Iou) >= 0.65: | |
save_file = os.path.join(pos_save_dir, "%s.jpg" % p_idx) | |
f1.write(save_file + ' 1 %.2f %.2f %.2f %.2f\n' % ( | |
offset_x1, offset_y1, offset_x2, offset_y2)) | |
cv2.imwrite(save_file, resized_im) | |
p_idx += 1 | |
elif np.max(Iou) >= 0.4: | |
save_file = os.path.join(part_save_dir, "%s.jpg" % d_idx) | |
f3.write(save_file + ' -1 %.2f %.2f %.2f %.2f\n' % ( | |
offset_x1, offset_y1, offset_x2, offset_y2)) | |
cv2.imwrite(save_file, resized_im) | |
d_idx += 1 | |
f1.close() | |
f2.close() | |
f3.close() | |
def model_store_path(): | |
return os.path.dirname(os.path.dirname(os.path.dirname(os.path.realpath(__file__))))+"/model_store" | |
def get_Onet_data(pnet_model, rnet_model): | |
gen_onet_data(traindata_store, annotation_file, pnet_model_file = pnet_model, rnet_model_file = rnet_model,prefix_path=prefix_path,use_cuda = True, vis = False) | |
def assembel_Onet_data(): | |
anno_list = [] | |
anno_list.append(onet_postive_file) | |
anno_list.append(onet_part_file) | |
anno_list.append(onet_neg_file) | |
anno_list.append(onet_landmark_file) | |
chose_count = assemble_data(imglist_filename_onet ,anno_list) | |
print("ONet train annotation result file path:%s" % imglist_filename_onet) | |
def gen_landmark_48(anno_file, data_dir, prefix = ''): | |
size = 48 | |
image_id = 0 | |
landmark_imgs_save_dir = os.path.join(data_dir,"48/landmark") | |
if not os.path.exists(landmark_imgs_save_dir): | |
os.makedirs(landmark_imgs_save_dir) | |
anno_dir = './anno_store' | |
if not os.path.exists(anno_dir): | |
os.makedirs(anno_dir) | |
landmark_anno_filename = "landmark_48.txt" | |
save_landmark_anno = os.path.join(anno_dir,landmark_anno_filename) | |
# print(save_landmark_anno) | |
# time.sleep(5) | |
f = open(save_landmark_anno, 'w') | |
# dstdir = "train_landmark_few" | |
with open(anno_file, 'r') as f2: | |
annotations = f2.readlines() | |
num = len(annotations) | |
print("%d total images" % num) | |
l_idx =0 | |
idx = 0 | |
# image_path bbox landmark(5*2) | |
for annotation in annotations: | |
# print imgPath | |
annotation = annotation.strip().split(' ') | |
assert len(annotation)==15,"each line should have 15 element" | |
im_path = os.path.join('./data_set/face_landmark/CNN_FacePoint/train/',annotation[0].replace("\\", "/")) | |
gt_box = list(map(float, annotation[1:5])) | |
# gt_box = [gt_box[0], gt_box[2], gt_box[1], gt_box[3]] | |
gt_box = np.array(gt_box, dtype=np.int32) | |
landmark = list(map(float, annotation[5:])) | |
landmark = np.array(landmark, dtype=np.float) | |
img = cv2.imread(im_path) | |
# print(im_path) | |
assert (img is not None) | |
height, width, channel = img.shape | |
# crop_face = img[gt_box[1]:gt_box[3]+1, gt_box[0]:gt_box[2]+1] | |
# crop_face = cv2.resize(crop_face,(size,size)) | |
idx = idx + 1 | |
if idx % 100 == 0: | |
print("%d images done, landmark images: %d"%(idx,l_idx)) | |
# print(im_path) | |
# print(gt_box) | |
x1, x2, y1, y2 = gt_box | |
gt_box[1] = y1 | |
gt_box[2] = x2 | |
# time.sleep(5) | |
# gt's width | |
w = x2 - x1 + 1 | |
# gt's height | |
h = y2 - y1 + 1 | |
if max(w, h) < 40 or x1 < 0 or y1 < 0: | |
continue | |
# random shift | |
for i in range(10): | |
bbox_size = np.random.randint(int(min(w, h) * 0.8), np.ceil(1.25 * max(w, h))) | |
delta_x = np.random.randint(-w * 0.2, w * 0.2) | |
delta_y = np.random.randint(-h * 0.2, h * 0.2) | |
nx1 = max(x1 + w / 2 - bbox_size / 2 + delta_x, 0) | |
ny1 = max(y1 + h / 2 - bbox_size / 2 + delta_y, 0) | |
nx2 = nx1 + bbox_size | |
ny2 = ny1 + bbox_size | |
if nx2 > width or ny2 > height: | |
continue | |
crop_box = np.array([nx1, ny1, nx2, ny2]) | |
cropped_im = img[int(ny1):int(ny2) + 1, int(nx1):int(nx2) + 1, :] | |
resized_im = cv2.resize(cropped_im, (size, size),interpolation=cv2.INTER_LINEAR) | |
offset_x1 = (x1 - nx1) / float(bbox_size) | |
offset_y1 = (y1 - ny1) / float(bbox_size) | |
offset_x2 = (x2 - nx2) / float(bbox_size) | |
offset_y2 = (y2 - ny2) / float(bbox_size) | |
offset_left_eye_x = (landmark[0] - nx1) / float(bbox_size) | |
offset_left_eye_y = (landmark[1] - ny1) / float(bbox_size) | |
offset_right_eye_x = (landmark[2] - nx1) / float(bbox_size) | |
offset_right_eye_y = (landmark[3] - ny1) / float(bbox_size) | |
offset_nose_x = (landmark[4] - nx1) / float(bbox_size) | |
offset_nose_y = (landmark[5] - ny1) / float(bbox_size) | |
offset_left_mouth_x = (landmark[6] - nx1) / float(bbox_size) | |
offset_left_mouth_y = (landmark[7] - ny1) / float(bbox_size) | |
offset_right_mouth_x = (landmark[8] - nx1) / float(bbox_size) | |
offset_right_mouth_y = (landmark[9] - ny1) / float(bbox_size) | |
# cal iou | |
iou = IoU(crop_box.astype(np.float), np.expand_dims(gt_box.astype(np.float), 0)) | |
# print(iou) | |
if iou > 0.65: | |
save_file = os.path.join(landmark_imgs_save_dir, "%s.jpg" % l_idx) | |
cv2.imwrite(save_file, resized_im) | |
f.write(save_file + ' -2 %.2f %.2f %.2f %.2f %.2f %.2f %.2f %.2f %.2f %.2f %.2f %.2f %.2f %.2f \n' % \ | |
(offset_x1, offset_y1, offset_x2, offset_y2, \ | |
offset_left_eye_x,offset_left_eye_y,offset_right_eye_x,offset_right_eye_y,offset_nose_x,offset_nose_y,offset_left_mouth_x,offset_left_mouth_y,offset_right_mouth_x,offset_right_mouth_y)) | |
# print(save_file) | |
# print(save_landmark_anno) | |
l_idx += 1 | |
f.close() | |
def parse_args(): | |
parser = argparse.ArgumentParser(description='Get data', | |
formatter_class=argparse.ArgumentDefaultsHelpFormatter) | |
parser.add_argument('--net', dest='net', help='which net to show', type=str) | |
parser.add_argument('--pnet_path', default="./model_store/pnet_epoch_20.pt",help='path to pnet model', type=str) | |
parser.add_argument('--rnet_path', default="./model_store/rnet_epoch_20.pt",help='path to rnet model', type=str) | |
parser.add_argument('--use_cuda', default=True,help='use cuda', type=bool) | |
args = parser.parse_args() | |
return args | |
#-----------------------------------------------------------------------------------------------------------------------------------------------# | |
if __name__ == '__main__': | |
args = parse_args() | |
dir = 'anno_store' | |
if not os.path.exists(dir): | |
os.makedirs(dir) | |
if args.net == "pnet": | |
wider_face(txt_from_path, anno_file) | |
get_Pnet_data() | |
assembel_Pnet_data() | |
elif args.net == "rnet": | |
get_Rnet_data(args.pnet_path) | |
assembel_Rnet_data() | |
elif args.net == "onet": | |
get_Onet_data(args.pnet_path, args.rnet_path) | |
gen_landmark_48(annotation_file_lm, traindata_store, prefix_path_lm) | |
assembel_Onet_data() |