scripts: use data config defined class names
Shorten name of --data-config-path argument to --data-config
This commit is contained in:
parent
89daa407e5
commit
9c0c1f23ab
10
detect.py
10
detect.py
|
@ -10,6 +10,7 @@ from utils import torch_utils
|
||||||
|
|
||||||
def detect(
|
def detect(
|
||||||
net_config_path,
|
net_config_path,
|
||||||
|
data_config_path,
|
||||||
images_path,
|
images_path,
|
||||||
weights_file_path='weights/yolov3.pt',
|
weights_file_path='weights/yolov3.pt',
|
||||||
output='output',
|
output='output',
|
||||||
|
@ -19,7 +20,6 @@ def detect(
|
||||||
nms_thres=0.45,
|
nms_thres=0.45,
|
||||||
save_txt=False,
|
save_txt=False,
|
||||||
save_images=False,
|
save_images=False,
|
||||||
class_path='data/coco.names',
|
|
||||||
):
|
):
|
||||||
|
|
||||||
device = torch_utils.select_device()
|
device = torch_utils.select_device()
|
||||||
|
@ -28,6 +28,8 @@ def detect(
|
||||||
os.system('rm -rf ' + output)
|
os.system('rm -rf ' + output)
|
||||||
os.makedirs(output, exist_ok=True)
|
os.makedirs(output, exist_ok=True)
|
||||||
|
|
||||||
|
data_config = parse_data_config(data_config_path)
|
||||||
|
|
||||||
# Load model
|
# Load model
|
||||||
model = Darknet(net_config_path, img_size)
|
model = Darknet(net_config_path, img_size)
|
||||||
|
|
||||||
|
@ -55,7 +57,7 @@ def detect(
|
||||||
model.to(device).eval()
|
model.to(device).eval()
|
||||||
|
|
||||||
# Set Dataloader
|
# Set Dataloader
|
||||||
classes = load_classes(class_path) # Extracts class labels from file
|
classes = load_classes(data_config['names']) # Extracts class labels from file
|
||||||
dataloader = load_images(images_path, batch_size=batch_size, img_size=img_size)
|
dataloader = load_images(images_path, batch_size=batch_size, img_size=img_size)
|
||||||
|
|
||||||
imgs = [] # Stores image paths
|
imgs = [] # Stores image paths
|
||||||
|
@ -151,7 +153,7 @@ if __name__ == '__main__':
|
||||||
parser.add_argument('--txt-out', type=bool, default=False)
|
parser.add_argument('--txt-out', type=bool, default=False)
|
||||||
|
|
||||||
parser.add_argument('--cfg', type=str, default='cfg/yolov3.cfg', help='cfg file path')
|
parser.add_argument('--cfg', type=str, default='cfg/yolov3.cfg', help='cfg file path')
|
||||||
parser.add_argument('--class-path', type=str, default='data/coco.names', help='path to class label file')
|
parser.add_argument('--data-config', type=str, default='cfg/coco.data', help='path to data config file')
|
||||||
parser.add_argument('--conf-thres', type=float, default=0.50, help='object confidence threshold')
|
parser.add_argument('--conf-thres', type=float, default=0.50, help='object confidence threshold')
|
||||||
parser.add_argument('--nms-thres', type=float, default=0.45, help='iou threshold for non-maximum suppression')
|
parser.add_argument('--nms-thres', type=float, default=0.45, help='iou threshold for non-maximum suppression')
|
||||||
parser.add_argument('--batch-size', type=int, default=1, help='size of the batches')
|
parser.add_argument('--batch-size', type=int, default=1, help='size of the batches')
|
||||||
|
@ -165,6 +167,7 @@ if __name__ == '__main__':
|
||||||
|
|
||||||
detect(
|
detect(
|
||||||
opt.cfg,
|
opt.cfg,
|
||||||
|
opt.data_config,
|
||||||
opt.image_folder,
|
opt.image_folder,
|
||||||
output=opt.output_folder,
|
output=opt.output_folder,
|
||||||
batch_size=opt.batch_size,
|
batch_size=opt.batch_size,
|
||||||
|
@ -173,5 +176,4 @@ if __name__ == '__main__':
|
||||||
nms_thres=opt.nms_thres,
|
nms_thres=opt.nms_thres,
|
||||||
save_txt=opt.txt_out,
|
save_txt=opt.txt_out,
|
||||||
save_images=opt.plot_flag,
|
save_images=opt.plot_flag,
|
||||||
class_path=opt.class_path,
|
|
||||||
)
|
)
|
||||||
|
|
10
test.py
10
test.py
|
@ -11,7 +11,6 @@ def test(
|
||||||
net_config_path,
|
net_config_path,
|
||||||
data_config_path,
|
data_config_path,
|
||||||
weights_file_path,
|
weights_file_path,
|
||||||
class_path=None,
|
|
||||||
batch_size=16,
|
batch_size=16,
|
||||||
img_size=416,
|
img_size=416,
|
||||||
iou_thres=0.5,
|
iou_thres=0.5,
|
||||||
|
@ -118,8 +117,7 @@ def test(
|
||||||
# Print mAP per class
|
# Print mAP per class
|
||||||
print('%11s' * 5 % ('Image', 'Total', 'P', 'R', 'mAP') + '\n\nmAP Per Class:')
|
print('%11s' * 5 % ('Image', 'Total', 'P', 'R', 'mAP') + '\n\nmAP Per Class:')
|
||||||
|
|
||||||
if class_path:
|
classes = load_classes(data_config['names']) # Extracts class labels from file
|
||||||
classes = load_classes(class_path) # Extracts class labels from file
|
|
||||||
for i, c in enumerate(classes):
|
for i, c in enumerate(classes):
|
||||||
print('%15s: %-.4f' % (c, AP_accum[i] / AP_accum_count[i]))
|
print('%15s: %-.4f' % (c, AP_accum[i] / AP_accum_count[i]))
|
||||||
|
|
||||||
|
@ -132,9 +130,8 @@ if __name__ == '__main__':
|
||||||
parser = argparse.ArgumentParser(prog='test.py')
|
parser = argparse.ArgumentParser(prog='test.py')
|
||||||
parser.add_argument('--batch-size', type=int, default=32, help='size of each image batch')
|
parser.add_argument('--batch-size', type=int, default=32, help='size of each image batch')
|
||||||
parser.add_argument('--cfg', type=str, default='cfg/yolov3.cfg', help='path to model config file')
|
parser.add_argument('--cfg', type=str, default='cfg/yolov3.cfg', help='path to model config file')
|
||||||
parser.add_argument('--data-config-path', type=str, default='cfg/coco.data', help='path to data config file')
|
parser.add_argument('--data-config', type=str, default='cfg/coco.data', help='path to data config file')
|
||||||
parser.add_argument('--weights', type=str, default='weights/yolov3.pt', help='path to weights file')
|
parser.add_argument('--weights', type=str, default='weights/yolov3.pt', help='path to weights file')
|
||||||
parser.add_argument('--class-path', type=str, default='data/coco.names', help='path to class label file')
|
|
||||||
parser.add_argument('--iou-thres', type=float, default=0.5, help='iou threshold required to qualify as detected')
|
parser.add_argument('--iou-thres', type=float, default=0.5, help='iou threshold required to qualify as detected')
|
||||||
parser.add_argument('--conf-thres', type=float, default=0.3, help='object confidence threshold')
|
parser.add_argument('--conf-thres', type=float, default=0.3, help='object confidence threshold')
|
||||||
parser.add_argument('--nms-thres', type=float, default=0.45, help='iou threshold for non-maximum suppression')
|
parser.add_argument('--nms-thres', type=float, default=0.45, help='iou threshold for non-maximum suppression')
|
||||||
|
@ -147,9 +144,8 @@ if __name__ == '__main__':
|
||||||
|
|
||||||
mAP = test(
|
mAP = test(
|
||||||
opt.cfg,
|
opt.cfg,
|
||||||
opt.data_config_path,
|
opt.data_config,
|
||||||
opt.weights,
|
opt.weights,
|
||||||
class_path=opt.class_path,
|
|
||||||
batch_size=opt.batch_size,
|
batch_size=opt.batch_size,
|
||||||
img_size=opt.img_size,
|
img_size=opt.img_size,
|
||||||
iou_thres=opt.iou_thres,
|
iou_thres=opt.iou_thres,
|
||||||
|
|
4
train.py
4
train.py
|
@ -220,7 +220,7 @@ if __name__ == '__main__':
|
||||||
parser = argparse.ArgumentParser()
|
parser = argparse.ArgumentParser()
|
||||||
parser.add_argument('--epochs', type=int, default=100, help='number of epochs')
|
parser.add_argument('--epochs', type=int, default=100, help='number of epochs')
|
||||||
parser.add_argument('--batch-size', type=int, default=16, help='size of each image batch')
|
parser.add_argument('--batch-size', type=int, default=16, help='size of each image batch')
|
||||||
parser.add_argument('--data-config-path', type=str, default='cfg/coco.data', help='data config file path')
|
parser.add_argument('--data-config', type=str, default='cfg/coco.data', help='path to data config file')
|
||||||
parser.add_argument('--cfg', type=str, default='cfg/yolov3.cfg', help='cfg file path')
|
parser.add_argument('--cfg', type=str, default='cfg/yolov3.cfg', help='cfg file path')
|
||||||
parser.add_argument('--multi-scale', default=False, help='random image sizes per batch 320 - 608')
|
parser.add_argument('--multi-scale', default=False, help='random image sizes per batch 320 - 608')
|
||||||
parser.add_argument('--img-size', type=int, default=32 * 13, help='pixels')
|
parser.add_argument('--img-size', type=int, default=32 * 13, help='pixels')
|
||||||
|
@ -236,7 +236,7 @@ if __name__ == '__main__':
|
||||||
torch.cuda.empty_cache()
|
torch.cuda.empty_cache()
|
||||||
train(
|
train(
|
||||||
opt.cfg,
|
opt.cfg,
|
||||||
opt.data_config_path,
|
opt.data_config,
|
||||||
img_size=opt.img_size,
|
img_size=opt.img_size,
|
||||||
resume=opt.resume,
|
resume=opt.resume,
|
||||||
epochs=opt.epochs,
|
epochs=opt.epochs,
|
||||||
|
|
Loading…
Reference in New Issue