(args)
| 157 | |
| 158 | |
| 159 | def main(args): |
| 160 | output_dir = os.path.realpath(args.output) |
| 161 | os.makedirs(output_dir, exist_ok=True) |
| 162 | |
| 163 | labels = ["person","bicycle","car","motorcycle","airplane","bus","train","truck","boat","traffic light","fire hydrant","stop sign","parking meter","bench","bird","cat","dog","horse","sheep","cow","elephant","bear","zebra","giraffe","backpack","umbrella","handbag","tie","suitcase","frisbee","skis","snowboard","sports ball","kite","baseball bat","baseball glove","skateboard","surfboard","tennis racket","bottle","wine glass","cup","fork","knife","spoon","bowl","banana","apple","sandwich","orange","broccoli","carrot","hot dog","pizza","donut","cake","chair","couch","potted plant","bed","dining table","toilet","tv","laptop","mouse","remote","keyboard","cell phone","microwave","oven","toaster","sink","refrigerator","book","clock","vase","scissors","teddy bear","hair drier", "toothbrush"] |
| 164 | |
| 165 | trt_infer = TensorRTInfer(args.engine) |
| 166 | batcher = ImageBatcher(args.input, *trt_infer.input_spec(), config_file=args.det2_config) |
| 167 | for batch, images, scales in batcher.get_batch(): |
| 168 | print("Processing Image {} / {}".format(batcher.image_index, batcher.num_images), end="\r") |
| 169 | detections = trt_infer.infer(batch, scales, args.nms_threshold) |
| 170 | for i in range(len(images)): |
| 171 | basename = os.path.splitext(os.path.basename(images[i]))[0] |
| 172 | # Image Visualizations |
| 173 | output_path = os.path.join(output_dir, "{}.png".format(basename)) |
| 174 | visualize_detections(images[i], output_path, detections[i], labels, args.iou_threshold) |
| 175 | # Text Results |
| 176 | output_results = "" |
| 177 | for d in detections[i]: |
| 178 | line = [d['xmin'], d['ymin'], d['xmax'], d['ymax'], d['score'], d['class']] |
| 179 | output_results += "\t".join([str(f) for f in line]) + "\n" |
| 180 | with open(os.path.join(args.output, "{}.txt".format(basename)), "w") as f: |
| 181 | f.write(output_results) |
| 182 | print() |
| 183 | print("Finished Processing") |
| 184 | |
| 185 | |
| 186 | if __name__ == "__main__": |
no test coverage detected