Get image path list in TEST mode
(infer_dir, infer_img)
| 1112 | |
| 1113 | |
| 1114 | def get_test_images(infer_dir, infer_img): |
| 1115 | """ |
| 1116 | Get image path list in TEST mode |
| 1117 | """ |
| 1118 | assert infer_img is not None or infer_dir is not None, \ |
| 1119 | "--image_file or --image_dir should be set" |
| 1120 | assert infer_img is None or os.path.isfile(infer_img), \ |
| 1121 | "{} is not a file".format(infer_img) |
| 1122 | assert infer_dir is None or os.path.isdir(infer_dir), \ |
| 1123 | "{} is not a directory".format(infer_dir) |
| 1124 | |
| 1125 | # infer_img has a higher priority |
| 1126 | if infer_img and os.path.isfile(infer_img): |
| 1127 | return [infer_img] |
| 1128 | |
| 1129 | images = set() |
| 1130 | infer_dir = os.path.abspath(infer_dir) |
| 1131 | assert os.path.isdir(infer_dir), \ |
| 1132 | "infer_dir {} is not a directory".format(infer_dir) |
| 1133 | exts = ['jpg', 'jpeg', 'png', 'bmp'] |
| 1134 | exts += [ext.upper() for ext in exts] |
| 1135 | for ext in exts: |
| 1136 | images.update(glob.glob('{}/*.{}'.format(infer_dir, ext))) |
| 1137 | images = list(images) |
| 1138 | |
| 1139 | assert len(images) > 0, "no image found in {}".format(infer_dir) |
| 1140 | print("Found {} inference images in total.".format(len(images))) |
| 1141 | |
| 1142 | return images |
| 1143 | |
| 1144 | |
| 1145 | def visualize(image_list, result, labels, output_dir='output/', threshold=0.5): |