File size: 3,902 Bytes
901e379 |
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 |
"""
This code is used to batch detect images in a folder.
"""
import argparse
import os
import sys
import cv2
from vision.ssd.config.fd_config import define_img_size
parser = argparse.ArgumentParser(description='detect_imgs')
parser.add_argument('--net_type', default="RFB", type=str,
help='The network architecture ,optional: RFB (higher precision) or slim (faster)')
parser.add_argument('--input_size', default=320, type=int,
help='define network input size,default optional value 128/160/320/480/640/1280')
parser.add_argument('--threshold', default=0.65, type=float,
help='score threshold')
parser.add_argument('--candidate_size', default=1500, type=int,
help='nms candidate size')
parser.add_argument('--path', default="D:/Database/face_detect/test/originalPics", type=str,
help='imgs dir')
parser.add_argument('--test_device', default="cpu", type=str,
help='cuda:0 or cpu')
args = parser.parse_args()
define_img_size(args.input_size) # must put define_img_size() before 'import create_mb_tiny_fd, create_mb_tiny_fd_predictor'
from vision.ssd.mb_tiny_fd import create_mb_tiny_fd, create_mb_tiny_fd_predictor
from vision.ssd.mb_tiny_RFB_fd import create_Mb_Tiny_RFB_fd, create_Mb_Tiny_RFB_fd_predictor
result_path = "./detect_imgs_results"
label_path = "./models/voc-model-labels.txt"
fd_result_path = 'D:/Database/face_detect/test/rfb_fd_result.txt'
fddb_txt_path = 'D:/Database/face_detect/test/FDDB-folds/FDDB-fold-01-10_2845.txt'
test_device = args.test_device
class_names = [name.strip() for name in open(label_path).readlines()]
if args.net_type == 'slim':
model_path = "models/pretrained/version-slim-320.pth"
net = create_mb_tiny_fd(len(class_names), is_test=True, device=test_device)
predictor = create_mb_tiny_fd_predictor(net, candidate_size=args.candidate_size, device=test_device)
elif args.net_type == 'RFB':
model_path = "models/pretrained/version-RFB-320.pth"
net = create_Mb_Tiny_RFB_fd(len(class_names), is_test=True, device=test_device)
predictor = create_Mb_Tiny_RFB_fd_predictor(net, candidate_size=args.candidate_size, device=test_device)
else:
print("The net type is wrong!")
sys.exit(1)
net.load(model_path)
def get_file_names(dir_path):
file_list = os.listdir(dir_path)
total_file_list = list()
for entry in file_list:
full_path = os.path.join(dir_path, entry)
if (os.path.isdir(full_path)):
total_file_list = total_file_list + get_file_names(full_path)
else:
total_file_list.append(full_path)
return total_file_list
def get_file_paths(txt_path):
path_list = list()
with open(txt_path, "r") as txt_file:
for line in txt_file:
path_list.append(line.strip())
return path_list
if __name__ == '__main__':
if not os.path.exists(result_path):
os.makedirs(result_path)
listdir = get_file_paths(fddb_txt_path)
total_count = 0
correct_count = 0
for file_path in listdir:
filename = file_path
img_path = os.path.join(args.path, filename)
orig_image = cv2.imread(img_path + ".jpg")
if orig_image is None:
continue
print("filename: ", filename)
image = cv2.cvtColor(orig_image, cv2.COLOR_BGR2RGB)
boxes, labels, probs = predictor.predict(image, args.candidate_size / 2, args.threshold)
with open(fd_result_path, "a") as fd_result_file:
print(filename, file=fd_result_file)
print(boxes.size(0), file=fd_result_file)
for i in range(boxes.size(0)):
box = boxes[i, :]
score = f"{probs[i]:.3f}"
print(f"{box[0]:.3f}", f"{box[1]:.3f}", f"{box[2] - box[0]:.3f}", f"{box[3] - box[1]:.3f}", score, file=fd_result_file)
|