41 lines
1.2 KiB
Python
41 lines
1.2 KiB
Python
|
#!/usr/bin/env python
|
||
|
|
|
||
|
|
import argparse
|
||
|
|
from detector import Detector
|
||
|
|
import tempfile
|
||
|
|
from pathlib import Path
|
||
|
|
|
||
|
|
parser = argparse.ArgumentParser(
|
||
|
|
prog = 'YOLO CLI',
|
||
|
|
description = 'YOLO Command Line Interface.')
|
||
|
|
|
||
|
|
parser.add_argument('-w', '--yolo_weights',
|
||
|
|
help="Weight file for YOLOv3 object detection.")
|
||
|
|
parser.add_argument('-c', '--yolo_config',
|
||
|
|
help="Configuration file for YOLOv3 object detection")
|
||
|
|
parser.add_argument('-l', '--yolo_labels',
|
||
|
|
help="File containing list of object labels for YOLOv3 object detection")
|
||
|
|
|
||
|
|
parser.add_argument('files', metavar='FILE', type=str, nargs='*');
|
||
|
|
|
||
|
|
args = parser.parse_args()
|
||
|
|
|
||
|
|
outgoing_dir = Path(tempfile.mkdtemp())
|
||
|
|
|
||
|
|
classes = []
|
||
|
|
with open(args.yolo_labels) as f:
|
||
|
|
classes = [line.strip() for line in f.readlines()]
|
||
|
|
|
||
|
|
detector = Detector(
|
||
|
|
args.yolo_weights,
|
||
|
|
args.yolo_config,
|
||
|
|
classes,
|
||
|
|
outgoing_dir)
|
||
|
|
|
||
|
|
for filename in args.files:
|
||
|
|
print(filename + ":")
|
||
|
|
result = detector.detect_objects(filename)
|
||
|
|
print(" OUTPUT: " + str(result.outfile))
|
||
|
|
for detection in result.detections:
|
||
|
|
print(" " + detection.label + " (" + str(detection.confidence) + ")")
|