diff --git a/README.md b/README.md index 0ddad037..cbb37de9 100755 --- a/README.md +++ b/README.md @@ -56,7 +56,7 @@ Checkpoints are saved in `/checkpoints` directory. Run `detect.py` to apply trai Run `test.py` to validate the official YOLOv3 weights `checkpoints/yolov3.weights` against the 5000 validation images. You should obtain a mAP of .581 using this repo (https://github.com/ultralytics/yolov3), compared to .579 as reported in darknet (https://arxiv.org/abs/1804.02767). -Run `test.py -weights_path checkpoints/latest.pt` to validate against the latest training checkpoint. +Run `test.py --weights checkpoints/latest.pt` to validate against the latest training checkpoint. # Contact diff --git a/detect.py b/detect.py index cc2ef727..a104e8a3 100755 --- a/detect.py +++ b/detect.py @@ -5,45 +5,39 @@ from models import * from utils.datasets import * from utils.utils import * -f_path = os.path.dirname(os.path.realpath(__file__)) + '/' - -parser = argparse.ArgumentParser() -# Get data configuration - -parser.add_argument('-image_folder', type=str, default='data/samples', help='path to images') -parser.add_argument('-output_folder', type=str, default='output', help='path to outputs') -parser.add_argument('-plot_flag', type=bool, default=True) -parser.add_argument('-txt_out', type=bool, default=False) - -parser.add_argument('-cfg', type=str, default=f_path + 'cfg/yolov3.cfg', help='cfg file path') -parser.add_argument('-class_path', type=str, default=f_path + 'data/coco.names', help='path to class label file') -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('-batch_size', type=int, default=1, help='size of the batches') -parser.add_argument('-img_size', type=int, default=32 * 13, help='size of each image dimension') -opt = parser.parse_args() -print(opt) +from utils import torch_utils -def main(opt): +def detect( + net_config_path, + images_path, + weights_file_path='weights/yolov3.pt', + output='output', + batch_size=16, + img_size=416, + conf_thres=0.3, + nms_thres=0.45, + save_txt=False, + save_images=False, + class_path='data/coco.names', +): device = torch_utils.select_device() print("Using device: \"{}\"".format(device)) - os.system('rm -rf ' + opt.output_folder) - os.makedirs(opt.output_folder, exist_ok=True) + os.system('rm -rf ' + output) + os.makedirs(output, exist_ok=True) # Load model - model = Darknet(opt.cfg, opt.img_size) + model = Darknet(net_config_path, img_size) - weights_path = f_path + 'weights/yolov3.pt' - if weights_path.endswith('.pt'): # pytorch format - if weights_path.endswith('weights/yolov3.pt') and not os.path.isfile(weights_path): - os.system('wget https://storage.googleapis.com/ultralytics/yolov3.pt -O ' + weights_path) + if weights_file_path.endswith('.pt'): # pytorch format + if weights_file_path.endswith('weights/yolov3.pt') and not os.path.isfile(weights_file_path): + os.system('wget https://storage.googleapis.com/ultralytics/yolov3.pt -O ' + weights_file_path) else: # darknet format - load_weights(model, weights_path) + load_weights(model, weights_file_path) - checkpoint = torch.load(weights_path, map_location='cpu') + checkpoint = torch.load(weights_file_path, map_location='cpu') model.load_state_dict(checkpoint['model']) del checkpoint @@ -61,8 +55,8 @@ def main(opt): model.to(device).eval() # Set Dataloader - classes = load_classes(opt.class_path) # Extracts class labels from file - dataloader = load_images(opt.image_folder, batch_size=opt.batch_size, img_size=opt.img_size) + classes = load_classes(class_path) # Extracts class labels from file + dataloader = load_images(images_path, batch_size=batch_size, img_size=img_size) imgs = [] # Stores image paths img_detections = [] # Stores detections for each image index @@ -73,10 +67,10 @@ def main(opt): # Get detections with torch.no_grad(): pred = model(torch.from_numpy(img).unsqueeze(0).to(device)) - pred = pred[pred[:, :, 4] > opt.conf_thres] + pred = pred[pred[:, :, 4] > conf_thres] if len(pred) > 0: - detections = non_max_suppression(pred.unsqueeze(0), opt.conf_thres, opt.nms_thres) + detections = non_max_suppression(pred.unsqueeze(0), conf_thres, nms_thres) img_detections.extend(detections) imgs.extend(img_paths) @@ -93,15 +87,15 @@ def main(opt): for img_i, (path, detections) in enumerate(zip(imgs, img_detections)): print("image %g: '%s'" % (img_i, path)) - if opt.plot_flag: + if save_images: img = cv2.imread(path) # The amount of padding that was added - pad_x = max(img.shape[0] - img.shape[1], 0) * (opt.img_size / max(img.shape)) - pad_y = max(img.shape[1] - img.shape[0], 0) * (opt.img_size / max(img.shape)) + pad_x = max(img.shape[0] - img.shape[1], 0) * (img_size / max(img.shape)) + pad_y = max(img.shape[1] - img.shape[0], 0) * (img_size / max(img.shape)) # Image height and width after padding is removed - unpad_h = opt.img_size - pad_y - unpad_w = opt.img_size - pad_x + unpad_h = img_size - pad_y + unpad_w = img_size - pad_x # Draw bounding boxes and labels of detections if detections is not None: @@ -109,7 +103,7 @@ def main(opt): bbox_colors = random.sample(color_list, len(unique_classes)) # write results to .txt file - results_img_path = os.path.join(opt.output_folder, path.split('/')[-1]) + results_img_path = os.path.join(output, path.split('/')[-1]) results_txt_path = results_img_path + '.txt' if os.path.isfile(results_txt_path): os.remove(results_txt_path) @@ -129,24 +123,55 @@ def main(opt): x1, y1, x2, y2 = max(x1, 0), max(y1, 0), max(x2, 0), max(y2, 0) # write to file - if opt.txt_out: + if save_txt: with open(results_txt_path, 'a') as file: file.write(('%g %g %g %g %g %g \n') % (x1, y1, x2, y2, cls_pred, cls_conf * conf)) - if opt.plot_flag: + if save_images: # Add the bbox to the plot label = '%s %.2f' % (classes[int(cls_pred)], conf) color = bbox_colors[int(np.where(unique_classes == int(cls_pred))[0])] plot_one_box([x1, y1, x2, y2], img, label=label, color=color) - if opt.plot_flag: + if save_images: # Save generated image with detections cv2.imwrite(results_img_path.replace('.bmp', '.jpg').replace('.tif', '.jpg'), img) if platform == 'darwin': # MacOS (local) - os.system('open ' + opt.output_folder) + os.system('open ' + output) if __name__ == '__main__': + parser = argparse.ArgumentParser() + # Get data configuration + + parser.add_argument('--image-folder', type=str, default='data/samples', help='path to images') + parser.add_argument('--output-folder', type=str, default='output', help='path to outputs') + parser.add_argument('--plot-flag', type=bool, default=True) + 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('--class-path', type=str, default='data/coco.names', help='path to class label file') + 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('--batch-size', type=int, default=1, help='size of the batches') + parser.add_argument('--img-size', type=int, default=32 * 13, help='size of each image dimension') + opt = parser.parse_args() + print(opt) + torch.cuda.empty_cache() - main(opt) + + init_seeds() + + detect( + opt.cfg, + opt.image_folder, + output=opt.output_folder, + batch_size=opt.batch_size, + img_size=opt.img_size, + conf_thres=opt.conf_thres, + nms_thres=opt.nms_thres, + save_txt=opt.txt_out, + save_images=opt.plot_flag, + class_path=opt.class_path, + ) diff --git a/test.py b/test.py index c0cf476b..4630a0ee 100644 --- a/test.py +++ b/test.py @@ -6,47 +6,44 @@ from utils.utils import * from utils import torch_utils -parser = argparse.ArgumentParser(prog='test.py') -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('-data_config_path', type=str, default='cfg/coco.data', help='path to data config file') -parser.add_argument('-weights_path', 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('-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('-n_cpu', type=int, default=0, help='number of cpu threads to use during batch generation') -parser.add_argument('-img_size', type=int, default=416, help='size of each image dimension') -opt = parser.parse_args() -print(opt, end='\n\n') - -def main(opt): +def test( + net_config_path, + data_config_path, + weights_file_path, + class_path=None, + batch_size=16, + img_size=416, + iou_thres=0.5, + conf_thres=0.3, + nms_thres=0.45, + n_cpus=0, +): device = torch_utils.select_device() print("Using device: \"{}\"".format(device)) # Configure run - data_config = parse_data_config(opt.data_config_path) + data_config = parse_data_config(data_config_path) nC = int(data_config['classes']) # number of classes (80 for COCO) test_path = data_config['valid'] # Initiate model - model = Darknet(opt.cfg, opt.img_size) + model = Darknet(net_config_path, img_size) # Load weights - if opt.weights_path.endswith('.pt'): # pytorch format - checkpoint = torch.load(opt.weights_path, map_location='cpu') + if weights_file_path.endswith('.pt'): # pytorch format + checkpoint = torch.load(weights_file_path, map_location='cpu') model.load_state_dict(checkpoint['model']) del checkpoint else: # darknet format - load_weights(model, opt.weights_path) + load_weights(model, weights_file_path) model.to(device).eval() # Get dataloader # dataset = load_images_with_labels(test_path) - # dataloader = torch.utils.data.DataLoader(dataset, batch_size=opt.batch_size, shuffle=False, num_workers=opt.n_cpu) - dataloader = load_images_and_labels(test_path, batch_size=opt.batch_size, img_size=opt.img_size) + # dataloader = torch.utils.data.DataLoader(dataset, batch_size=batch_size, shuffle=False, num_workers=n_cpus) + dataloader = load_images_and_labels(test_path, batch_size=batch_size, img_size=img_size) print('%11s' * 5 % ('Image', 'Total', 'P', 'R', 'mAP')) outputs, mAPs, mR, mP, TP, confidence, pred_class, target_class = [], [], [], [], [], [], [], [] @@ -55,7 +52,7 @@ def main(opt): with torch.no_grad(): output = model(imgs.to(device)) - output = non_max_suppression(output, conf_thres=opt.conf_thres, nms_thres=opt.nms_thres) + output = non_max_suppression(output, conf_thres=conf_thres, nms_thres=nms_thres) # Compute average precision for each sample for sample_i, (labels, detections) in enumerate(zip(targets, output)): @@ -80,7 +77,7 @@ def main(opt): target_cls = labels[:, 0] # Extract target boxes as (x1, y1, x2, y2) - target_boxes = xywh2xyxy(labels[:, 1:5]) * opt.img_size + target_boxes = xywh2xyxy(labels[:, 1:5]) * img_size detected = [] for *pred_bbox, conf, obj_conf, obj_pred in detections: @@ -91,7 +88,7 @@ def main(opt): # Extract index of largest overlap best_i = np.argmax(iou) # If overlap exceeds threshold and classification is correct mark as correct - if iou[best_i] > opt.iou_thres and obj_pred == labels[best_i, 0] and best_i not in detected: + if iou[best_i] > iou_thres and obj_pred == labels[best_i, 0] and best_i not in detected: correct.append(1) detected.append(best_i) else: @@ -121,9 +118,10 @@ def main(opt): # Print mAP per class print('%11s' * 5 % ('Image', 'Total', 'P', 'R', 'mAP') + '\n\nmAP Per Class:') - classes = load_classes(opt.class_path) # Extracts class labels from file - for i, c in enumerate(classes): - print('%15s: %-.4f' % (c, AP_accum[i] / AP_accum_count[i])) + if class_path: + classes = load_classes(class_path) # Extracts class labels from file + for i, c in enumerate(classes): + print('%15s: %-.4f' % (c, AP_accum[i] / AP_accum_count[i])) # Return mAP return mean_mAP, mean_R, mean_P @@ -131,6 +129,31 @@ def main(opt): if __name__ == '__main__': + parser = argparse.ArgumentParser(prog='test.py') + 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('--data-config-path', 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('--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('--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('--n-cpus', type=int, default=0, help='number of cpu threads to use during batch generation') + parser.add_argument('--img-size', type=int, default=416, help='size of each image dimension') + opt = parser.parse_args() + print(opt, end='\n\n') + init_seeds() - mAP = main(opt) + mAP = test( + opt.cfg, + opt.data_config_path, + opt.weights, + class_path=opt.class_path, + batch_size=opt.batch_size, + img_size=opt.img_size, + iou_thres=opt.iou_thres, + conf_thres=opt.conf_thres, + nms_thres=opt.nms_thres, + n_cpus=opt.n_cpus, + ) diff --git a/train.py b/train.py index 103d6f03..5bc9cf9f 100644 --- a/train.py +++ b/train.py @@ -8,51 +8,48 @@ from utils.utils import * from utils import torch_utils -parser = argparse.ArgumentParser() -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('-data_config_path', type=str, default='cfg/coco.data', help='data config 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('-img_size', type=int, default=32 * 13, help='pixels') -parser.add_argument('-resume', default=False, help='resume training flag') -parser.add_argument('-batch_report', default=False, help='report TP, FP, FN, P and R per batch (slower)') -parser.add_argument('-freeze_darknet53', default=False, help='freeze darknet53.conv.74 layers for first epoch') -parser.add_argument('-var', type=float, default=0, help='optional test variable') -opt = parser.parse_args() -if opt.multi_scale: # pass maximum multi_scale size - opt.img_size = 608 -print(opt) - # Import test.py to get mAP after each epoch -sys.argv[1:] = [] # delete any train.py command-line arguments before they reach test.py -import test # must follow sys.argv[1:] = [] +import test -def main(opt): +def train( + net_config_path, + data_config_path, + img_size=416, + resume=False, + epochs=100, + batch_size=16, + report=False, + multi_scale=False, + freeze_backbone=True, + var=0, +): device = torch_utils.select_device() print("Using device: \"{}\"".format(device)) - if not opt.multi_scale: + if not multi_scale: torch.backends.cudnn.benchmark = True os.makedirs('weights', exist_ok=True) # Configure run - data_config = parse_data_config(opt.data_config_path) + data_config = parse_data_config(data_config_path) num_classes = int(data_config['classes']) train_path = '../coco/trainvalno5k.txt' # Initialize model - model = Darknet(opt.cfg, opt.img_size) + model = Darknet(net_config_path, img_size) # Get dataloader - dataloader = load_images_and_labels(train_path, batch_size=opt.batch_size, img_size=opt.img_size, - multi_scale=opt.multi_scale, augment=True) + if multi_scale: # pass maximum multi_scale size + img_size = 608 + + dataloader = load_images_and_labels(train_path, batch_size=batch_size, img_size=img_size, + multi_scale=multi_scale, augment=True) lr0 = 0.001 - if opt.resume: + if resume: checkpoint = torch.load('weights/latest.pt', map_location='cpu') model.load_state_dict(checkpoint['model']) @@ -103,7 +100,7 @@ def main(opt): mean_recall, mean_precision = 0, 0 print('%11s' * 16 % ( 'Epoch', 'Batch', 'x', 'y', 'w', 'h', 'conf', 'cls', 'total', 'P', 'R', 'nTargets', 'TP', 'FP', 'FN', 'time')) - for epoch in range(opt.epochs): + for epoch in range(epochs): epoch += start_epoch # Update scheduler (automatic) @@ -118,7 +115,7 @@ def main(opt): g['lr'] = lr # Freeze darknet53.conv.74 layers for first epoch - if opt.freeze_darknet53: + if freeze_backbone is not False: if epoch == 0: for i, (name, p) in enumerate(model.named_parameters()): if int(name.split('.')[1]) < 75: # if layer < 75 @@ -143,7 +140,7 @@ def main(opt): g['lr'] = lr # Compute loss, compute gradient, update parameters - loss = model(imgs.to(device), targets, batch_report=opt.batch_report, var=opt.var) + loss = model(imgs.to(device), targets, batch_report=report, var=var) loss.backward() # accumulated_batches = 1 # accumulate gradient for 4 batches before stepping optimizer @@ -156,7 +153,7 @@ def main(opt): for key, val in model.losses.items(): rloss[key] = (rloss[key] * ui + val) / (ui + 1) - if opt.batch_report: + if report: TP, FP, FN = metrics metrics += model.losses['metrics'] @@ -173,7 +170,7 @@ def main(opt): mean_recall = recall[k].mean() s = ('%11s%11s' + '%11.3g' * 14) % ( - '%g/%g' % (epoch, opt.epochs - 1), '%g/%g' % (i, len(dataloader) - 1), rloss['x'], + '%g/%g' % (epoch, epochs - 1), '%g/%g' % (i, len(dataloader) - 1), rloss['x'], rloss['y'], rloss['w'], rloss['h'], rloss['conf'], rloss['cls'], rloss['loss'], mean_precision, mean_recall, model.losses['nT'], model.losses['TP'], model.losses['FP'], model.losses['FN'], time.time() - t1) @@ -201,8 +198,13 @@ def main(opt): os.system('cp weights/latest.pt weights/backup' + str(epoch) + '.pt') # Calculate mAP - test.opt.weights_path = 'weights/latest.pt' - mAP, R, P = test.main(test.opt) + mAP, R, P = test.test( + net_config_path, + data_config_path, + 'weights/latest.pt', + batch_size=batch_size, + img_size=img_size, + ) # Write epoch results with open('results.txt', 'a') as file: @@ -215,7 +217,32 @@ def main(opt): if __name__ == '__main__': + parser = argparse.ArgumentParser() + 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('--data-config-path', type=str, default='cfg/coco.data', help='data config 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('--img-size', type=int, default=32 * 13, help='pixels') + parser.add_argument('--resume', default=False, help='resume training flag') + parser.add_argument('--report', default=False, help='report TP, FP, FN, P and R per batch (slower)') + parser.add_argument('--freeze-darknet53', default=False, help='freeze darknet53.conv.74 layers for first epoch') + parser.add_argument('--var', type=float, default=0, help='optional test variable') + opt = parser.parse_args() + print(opt, end='\n\n') + init_seeds() torch.cuda.empty_cache() - main(opt) + train( + opt.cfg, + opt.data_config_path, + img_size=opt.img_size, + resume=opt.resume, + epochs=opt.epochs, + batch_size=opt.batch_size, + report=opt.report, + multi_scale=opt.multi_scale, + freeze_backbone=opt.freeze_darknet53, + var=opt.var, + ) diff --git a/utils/gcp.sh b/utils/gcp.sh old mode 100644 new mode 100755 index 3e27c169..130f3249 --- a/utils/gcp.sh +++ b/utils/gcp.sh @@ -4,28 +4,28 @@ sudo rm -rf yolov3 && git clone https://github.com/ultralytics/yolov3 && cd yolov3 && python3 train.py # Resume -python3 train.py -resume 1 +python3 train.py --resume 1 # Detect gsutil cp gs://ultralytics/yolov3.pt yolov3/weights python3 detect.py # Test -python3 test.py -img_size 416 -weights_path weights/latest.pt +python3 test.py --img_size 416 --weights weights/latest.pt # Test Darknet -python3 test.py -img_size 416 -weights_path ../darknet/backup/yolov3.backup +python3 test.py --img_size 416 --weights ../darknet/backup/yolov3.backup # Download and Test sudo rm -rf yolov3 && git clone https://github.com/ultralytics/yolov3 && cd yolov3 wget https://pjreddie.com/media/files/yolov3.weights -P weights -python3 test.py -img_size 416 -weights_path weights/backup5.pt -nms_thres 0.45 +python3 test.py --img_size 416 --weights weights/backup5.pt --nms_thres 0.45 # Download and Resume sudo rm -rf yolov3 && git clone https://github.com/ultralytics/yolov3 && cd yolov3 wget https://storage.googleapis.com/ultralytics/yolov3.pt -O weights/latest.pt -python3 train.py -img_size 416 -batch_size 16 -epochs 1 -resume 1 -python3 test.py -img_size 416 -weights_path weights/latest.pt -conf_thres 0.5 +python3 train.py --img_size 416 --batch_size 16 --epochs 1 --resume 1 +python3 test.py --img_size 416 --weights weights/latest.pt --conf_thres 0.5 # Copy latest.pt to bucket gsutil cp yolov3/weights/latest.pt gs://ultralytics @@ -36,6 +36,6 @@ wget https://storage.googleapis.com/ultralytics/latest.pt # Testing sudo rm -rf yolov3 && git clone https://github.com/ultralytics/yolov3 && cd yolov3 -python3 train.py -epochs 3 -var 64 +python3 train.py --epochs 3 --var 64 sudo shutdown