|
| 1 | +""" |
| 2 | +This code is used to batch detect images in a folder. |
| 3 | +""" |
| 4 | +import os, sys |
| 5 | +currentdir = os.path.dirname(os.path.realpath(__file__)) |
| 6 | +parentdir = os.path.dirname(currentdir) |
| 7 | +sys.path.append(parentdir) |
| 8 | + |
| 9 | +import argparse |
| 10 | +import os |
| 11 | +import sys |
| 12 | + |
| 13 | +import cv2 |
| 14 | + |
| 15 | +from vision.ssd.config.fd_config import define_img_size |
| 16 | + |
| 17 | +parser = argparse.ArgumentParser( |
| 18 | + description='detect_imgs') |
| 19 | + |
| 20 | +parser.add_argument('--net_type', default="RFB", type=str, |
| 21 | + help='The network architecture ,optional: RFB (higher precision) or slim (faster)') |
| 22 | +parser.add_argument('--input_size', default=1280, type=int, |
| 23 | + help='define network input size,default optional value 128/160/320/480/640/1280') |
| 24 | +parser.add_argument('--threshold', default=0.3, type=float, |
| 25 | + help='score threshold') |
| 26 | +parser.add_argument('--candidate_size', default=1200, type=int, |
| 27 | + help='nms candidate size') |
| 28 | +parser.add_argument('--path', default="imgs", type=str, |
| 29 | + help='imgs dir') |
| 30 | +parser.add_argument('--test_device', default="cpu", type=str, |
| 31 | + help='cuda:0 or cpu') |
| 32 | +args = parser.parse_args() |
| 33 | +define_img_size(args.input_size) # must put define_img_size() before 'import create_mb_tiny_fd, create_mb_tiny_fd_predictor' |
| 34 | + |
| 35 | +from vision.ssd.mb_tiny_fd import create_mb_tiny_fd, create_mb_tiny_fd_predictor |
| 36 | +from vision.ssd.mb_tiny_RFB_fd import create_Mb_Tiny_RFB_fd, create_Mb_Tiny_RFB_fd_predictor |
| 37 | + |
| 38 | +result_path = "detect_imgs_results" |
| 39 | +label_path = "./voc-model-labels.txt" |
| 40 | +test_device = args.test_device |
| 41 | + |
| 42 | +class_names = [name.strip() for name in open(label_path).readlines()] |
| 43 | + |
| 44 | +if args.net_type == 'RFB': |
| 45 | + model_path = "pretrained/RFB-640-masked_face-v2.pth" |
| 46 | + net = create_Mb_Tiny_RFB_fd(len(class_names), is_test=True, device=test_device) |
| 47 | + predictor = create_Mb_Tiny_RFB_fd_predictor(net, candidate_size=args.candidate_size, device=test_device) |
| 48 | +else: |
| 49 | + print("The net type is wrong!") |
| 50 | + sys.exit(1) |
| 51 | +net.load(model_path) |
| 52 | + |
| 53 | +if not os.path.exists(result_path): |
| 54 | + os.makedirs(result_path) |
| 55 | +listdir = os.listdir(args.path) |
| 56 | +sum = 0 |
| 57 | +for file_path in listdir: |
| 58 | + img_path = os.path.join(args.path, file_path) |
| 59 | + orig_image = cv2.imread(img_path) |
| 60 | + if orig_image is None: continue |
| 61 | + image = cv2.cvtColor(orig_image, cv2.COLOR_BGR2RGB) |
| 62 | + boxes, labels, probs = predictor.predict(image, args.candidate_size / 2, args.threshold) |
| 63 | + sum += boxes.size(0) |
| 64 | + for i in range(boxes.size(0)): |
| 65 | + box = boxes[i, :] |
| 66 | + label_index = labels[i].item() |
| 67 | + cv2.rectangle(orig_image, (int(box[0]), int(box[1])), (int(box[2]), int(box[3])), (0, 165, 255), 2) |
| 68 | + # label = f"""{voc_dataset.class_names[labels[i]]}: {probs[i]:.2f}""" |
| 69 | + label = f"{probs[i]:.2f}" |
| 70 | + # cv2.putText(orig_image, label, (box[0], box[1] - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 255), 2) |
| 71 | + cv2.putText(orig_image, class_names[label_index], |
| 72 | + (int(box[0]), int(box[1]) - 10), |
| 73 | + cv2.FONT_HERSHEY_SIMPLEX, |
| 74 | + 0.5, # font scale |
| 75 | + (100, 0, 255), |
| 76 | + 1) # line type |
| 77 | + cv2.putText(orig_image, str(boxes.size(0)), (30, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 255), 2) |
| 78 | + cv2.imwrite(os.path.join(result_path, file_path), orig_image) |
| 79 | + print(f"Found {len(probs)} faces. The output image is {result_path}") |
| 80 | +print(sum) |
0 commit comments