1. 程式人生 > >gluoncv 訓練自己的資料集,進行目標檢測

gluoncv 訓練自己的資料集,進行目標檢測

跑了一晚上的模型,實在佔GPU資源,這兩天已經有很多小朋友說我了。我選擇了其中一個引數。

https://github.com/dmlc/gluon-cv/blob/master/scripts/detection/faster_rcnn/train_faster_rcnn.py

train_faster_rcnn的修改之前就弄好了,這裡貼一個完整的。

"""Train Faster-RCNN end to end."""
import argparse
import os
# disable autotune
os.environ['MXNET_CUDNN_AUTOTUNE_DEFAULT
'] = '0' import logging import time import numpy as np import mxnet as mx from mxnet import nd from mxnet import gluon from mxnet import autograd import gluoncv as gcv from gluoncv import data as gdata from gluoncv import utils as gutils from gluoncv.model_zoo import get_model from gluoncv.data import
batchify from gluoncv.data.transforms.presets.rcnn import FasterRCNNDefaultTrainTransform from gluoncv.data.transforms.presets.rcnn import FasterRCNNDefaultValTransform from gluoncv.utils.metrics.voc_detection import VOC07MApMetric from gluoncv.utils.metrics.coco_detection import COCODetectionMetric
from gluoncv.utils.metrics.accuracy import Accuracy # add_lst from gluoncv.data import LstDetection def parse_args(): parser = argparse.ArgumentParser(description='Train Faster-RCNN networks e2e.') parser.add_argument('--network', type=str, default='resnet50_v1b', help="Base network name which serves as feature extraction base.") parser.add_argument('--dataset', type=str, default='voc', help='Training dataset. Now support voc and coco.') parser.add_argument('--num-workers', '-j', dest='num_workers', type=int, default=4, help='Number of data workers, you can use larger ' 'number to accelerate data loading, if you CPU and GPUs are powerful.') parser.add_argument('--gpus', type=str, default='0', help='Training with GPUs, you can specify 1,3 for example.') parser.add_argument('--epochs', type=str, default='', help='Training epochs.') parser.add_argument('--resume', type=str, default='', help='Resume from previously saved parameters if not None. ' 'For example, you can resume from ./faster_rcnn_xxx_0123.params') parser.add_argument('--start-epoch', type=int, default=0, help='Starting epoch for resuming, default is 0 for new training.' 'You can specify it to 100 for example to start from 100 epoch.') parser.add_argument('--lr', type=str, default='', help='Learning rate, default is 0.001 for voc single gpu training.') parser.add_argument('--lr-decay', type=float, default=0.1, help='decay rate of learning rate. default is 0.1.') parser.add_argument('--lr-decay-epoch', type=str, default='', help='epoches at which learning rate decays. default is 14,20 for voc.') parser.add_argument('--lr-warmup', type=str, default='', help='warmup iterations to adjust learning rate, default is 0 for voc.') parser.add_argument('--momentum', type=float, default=0.9, help='SGD momentum, default is 0.9') parser.add_argument('--wd', type=str, default='', help='Weight decay, default is 5e-4 for voc') parser.add_argument('--log-interval', type=int, default=100, help='Logging mini-batch interval. Default is 100.') parser.add_argument('--save-prefix', type=str, default='', help='Saving parameter prefix') parser.add_argument('--save-interval', type=int, default=1, help='Saving parameters epoch interval, best model will always be saved.') parser.add_argument('--val-interval', type=int, default=1, help='Epoch interval for validation, increase the number will reduce the ' 'training time if validation is slow.') parser.add_argument('--seed', type=int, default=233, help='Random seed to be fixed.') parser.add_argument('--verbose', dest='verbose', action='store_true', help='Print helpful debugging info once set.') parser.add_argument('--mixup', action='store_true', help='Use mixup training.') parser.add_argument('--no-mixup-epochs', type=int, default=20, help='Disable mixup training if enabled in the last N epochs.') args = parser.parse_args() if args.dataset == 'voc' or args.dataset == 'pedestrian': args.epochs = int(args.epochs) if args.epochs else 20 args.lr_decay_epoch = args.lr_decay_epoch if args.lr_decay_epoch else '14,20' args.lr = float(args.lr) if args.lr else 0.001 args.lr_warmup = args.lr_warmup if args.lr_warmup else -1 args.wd = float(args.wd) if args.wd else 5e-4 elif args.dataset == 'coco': args.epochs = int(args.epochs) if args.epochs else 26 args.lr_decay_epoch = args.lr_decay_epoch if args.lr_decay_epoch else '17,23' args.lr = float(args.lr) if args.lr else 0.00125 args.lr_warmup = args.lr_warmup if args.lr_warmup else 8000 args.wd = float(args.wd) if args.wd else 1e-4 num_gpus = len(args.gpus.split(',')) if num_gpus == 1: args.lr_warmup = -1 else: args.lr *= num_gpus args.lr_warmup /= num_gpus return args class RPNAccMetric(mx.metric.EvalMetric): def __init__(self): super(RPNAccMetric, self).__init__('RPNAcc') def update(self, labels, preds): # label: [rpn_label, rpn_weight] # preds: [rpn_cls_logits] rpn_label, rpn_weight = labels rpn_cls_logits = preds[0] # calculate num_inst (average on those fg anchors) num_inst = mx.nd.sum(rpn_weight) # cls_logits (b, c, h, w) red_label (b, 1, h, w) # pred_label = mx.nd.argmax(rpn_cls_logits, axis=1, keepdims=True) pred_label = mx.nd.sigmoid(rpn_cls_logits) >= 0.5 # label (b, 1, h, w) num_acc = mx.nd.sum((pred_label == rpn_label) * rpn_weight) self.sum_metric += num_acc.asscalar() self.num_inst += num_inst.asscalar() class RPNL1LossMetric(mx.metric.EvalMetric): def __init__(self): super(RPNL1LossMetric, self).__init__('RPNL1Loss') def update(self, labels, preds): # label = [rpn_bbox_target, rpn_bbox_weight] # pred = [rpn_bbox_reg] rpn_bbox_target, rpn_bbox_weight = labels rpn_bbox_reg = preds[0] # calculate num_inst (average on those fg anchors) num_inst = mx.nd.sum(rpn_bbox_weight) / 4 # calculate smooth_l1 loss = mx.nd.sum(rpn_bbox_weight * mx.nd.smooth_l1(rpn_bbox_reg - rpn_bbox_target, scalar=3)) self.sum_metric += loss.asscalar() self.num_inst += num_inst.asscalar() class RCNNAccMetric(mx.metric.EvalMetric): def __init__(self): super(RCNNAccMetric, self).__init__('RCNNAcc') def update(self, labels, preds): # label = [rcnn_label] # pred = [rcnn_cls] rcnn_label = labels[0] rcnn_cls = preds[0] # calculate num_acc pred_label = mx.nd.argmax(rcnn_cls, axis=-1) num_acc = mx.nd.sum(pred_label == rcnn_label) self.sum_metric += num_acc.asscalar() self.num_inst += rcnn_label.size class RCNNL1LossMetric(mx.metric.EvalMetric): def __init__(self): super(RCNNL1LossMetric, self).__init__('RCNNL1Loss') def update(self, labels, preds): # label = [rcnn_bbox_target, rcnn_bbox_weight] # pred = [rcnn_reg] rcnn_bbox_target, rcnn_bbox_weight = labels rcnn_bbox_reg = preds[0] # calculate num_inst num_inst = mx.nd.sum(rcnn_bbox_weight) / 4 # calculate smooth_l1 loss = mx.nd.sum(rcnn_bbox_weight * mx.nd.smooth_l1(rcnn_bbox_reg - rcnn_bbox_target, scalar=1)) self.sum_metric += loss.asscalar() self.num_inst += num_inst.asscalar() def get_dataset(dataset, args): if dataset.lower() == 'voc': train_dataset = gdata.VOCDetection( splits=[(2007, 'trainval'), (2012, 'trainval')]) val_dataset = gdata.VOCDetection( splits=[(2007, 'test')]) #print(val_dataset.classes) #('aeroplane', 'bicycle', 'bird', 'boat', 'bottle', 'bus', 'car', 'cat', 'chair', 'cow', 'diningtable', 'dog', 'horse', 'motorbike', 'person', 'pottedplant', 'sheep', 'sofa', 'train', 'tvmonitor') val_metric = VOC07MApMetric(iou_thresh=0.5, class_names=val_dataset.classes) elif dataset.lower() == 'coco': train_dataset = gdata.COCODetection(splits='instances_train2017', use_crowd=False) val_dataset = gdata.COCODetection(splits='instances_val2017', skip_empty=False) val_metric = COCODetectionMetric(val_dataset, args.save_prefix + '_eval', cleanup=True) elif dataset.lower() == 'pedestrian': lst_dataset = LstDetection('train_val.lst',root=os.path.expanduser('.')) print(len(lst_dataset)) first_img = lst_dataset[0][0] print(first_img.shape) print(lst_dataset[0][1]) train_dataset = LstDetection('train.lst',root=os.path.expanduser('.')) val_dataset = LstDetection('val.lst',root=os.path.expanduser('.')) classs = ('pedestrian',) val_metric = VOC07MApMetric(iou_thresh=0.5,class_names=classs) else: raise NotImplementedError('Dataset: {} not implemented.'.format(dataset)) if args.mixup: from gluoncv.data.mixup import MixupDetection train_dataset = MixupDetection(train_dataset) return train_dataset, val_dataset, val_metric def get_dataloader(net, train_dataset, val_dataset, batch_size, num_workers): """Get dataloader.""" train_bfn = batchify.Tuple(*[batchify.Append() for _ in range(5)]) train_loader = mx.gluon.data.DataLoader( train_dataset.transform(FasterRCNNDefaultTrainTransform(net.short, net.max_size, net)), batch_size, True, batchify_fn=train_bfn, last_batch='rollover', num_workers=num_workers) val_bfn = batchify.Tuple(*[batchify.Append() for _ in range(3)]) val_loader = mx.gluon.data.DataLoader( val_dataset.transform(FasterRCNNDefaultValTransform(net.short, net.max_size)), batch_size, False, batchify_fn=val_bfn, last_batch='keep', num_workers=num_workers) return train_loader, val_loader def save_params(net, logger, best_map, current_map, epoch, save_interval, prefix): current_map = float(current_map) if current_map > best_map[0]: logger.info('[Epoch {}] mAP {} higher than current best {} saving to {}'.format( epoch, current_map, best_map, '{:s}_best.params'.format(prefix))) best_map[0] = current_map net.save_parameters('{:s}_best.params'.format(prefix)) with open(prefix+'_best_map.log', 'a') as f: f.write('{:04d}:\t{:.4f}\n'.format(epoch, current_map)) if save_interval and (epoch + 1) % save_interval == 0: logger.info('[Epoch {}] Saving parameters to {}'.format( epoch, '{:s}_{:04d}_{:.4f}.params'.format(prefix, epoch, current_map))) net.save_parameters('{:s}_{:04d}_{:.4f}.params'.format(prefix, epoch, current_map)) def split_and_load(batch, ctx_list): """Split data to 1 batch each device.""" num_ctx = len(ctx_list) new_batch = [] for i, data in enumerate(batch): new_data = [x.as_in_context(ctx) for x, ctx in zip(data, ctx_list)] new_batch.append(new_data) return new_batch def validate(net, val_data, ctx, eval_metric): """Test on validation dataset.""" clipper = gcv.nn.bbox.BBoxClipToImage() eval_metric.reset() net.hybridize(static_alloc=True) for batch in val_data: batch = split_and_load(batch, ctx_list=ctx) det_bboxes = [] det_ids = [] det_scores = [] gt_bboxes = [] gt_ids = [] gt_difficults = [] for x, y, im_scale in zip(*batch): # get prediction results ids, scores, bboxes = net(x) det_ids.append(ids) det_scores.append(scores) # clip to image size det_bboxes.append(clipper(bboxes, x)) # rescale to original resolution im_scale = im_scale.reshape((-1)).asscalar() det_bboxes[-1] *= im_scale # split ground truths gt_ids.append(y.slice_axis(axis=-1, begin=4, end=5)) gt_bboxes.append(y.slice_axis(axis=-1, begin=0, end=4)) gt_bboxes[-1] *= im_scale gt_difficults.append(y.slice_axis(axis=-1, begin=5, end=6) if y.shape[-1] > 5 else None) # update metric for det_bbox, det_id, det_score, gt_bbox, gt_id, gt_diff in zip(det_bboxes, det_ids, det_scores, gt_bboxes, gt_ids, gt_difficults): eval_metric.update(det_bbox, det_id, det_score, gt_bbox, gt_id, gt_diff) return eval_metric.get() def get_lr_at_iter(alpha): return 1. / 3. * (1 - alpha) + alpha def train(net, train_data, val_data, eval_metric, ctx, args): """Training pipeline""" net.collect_params().setattr('grad_req', 'null') net.collect_train_params().setattr('grad_req', 'write') trainer = gluon.Trainer( net.collect_train_params(), # fix batchnorm, fix first stage, etc... 'sgd', {'learning_rate': args.lr, 'wd': args.wd, 'momentum': args.momentum, 'clip_gradient': 5}) # lr decay policy lr_decay = float(args.lr_decay) lr_steps = sorted([float(ls) for ls in args.lr_decay_epoch.split(',') if ls.strip()]) lr_warmup = float(args.lr_warmup) # avoid int division # TODO(zhreshold) losses? rpn_cls_loss = mx.gluon.loss.SigmoidBinaryCrossEntropyLoss(from_sigmoid=False) rpn_box_loss = mx.gluon.loss.HuberLoss(rho=1/9.) # == smoothl1 rcnn_cls_loss = mx.gluon.loss.SoftmaxCrossEntropyLoss() rcnn_box_loss = mx.gluon.loss.HuberLoss() # == smoothl1 metrics = [mx.metric.Loss('RPN_Conf'), mx.metric.Loss('RPN_SmoothL1'), mx.metric.Loss('RCNN_CrossEntropy'), mx.metric.Loss('RCNN_SmoothL1'),] rpn_acc_metric = RPNAccMetric() rpn_bbox_metric = RPNL1LossMetric() rcnn_acc_metric = RCNNAccMetric() rcnn_bbox_metric = RCNNL1LossMetric() metrics2 = [rpn_acc_metric, rpn_bbox_metric, rcnn_acc_metric, rcnn_bbox_metric] # set up logger logging.basicConfig() logger = logging.getLogger() logger.setLevel(logging.INFO) log_file_path = args.save_prefix + '_train.log' log_dir = os.path.dirname(log_file_path) if log_dir and not os.path.exists(log_dir): os.makedirs(log_dir) fh = logging.FileHandler(log_file_path) logger.addHandler(fh) logger.info(args) if args.verbose: logger.info('Trainable parameters:') logger.info(net.collect_train_params().keys()) logger.info('Start training from [Epoch {}]'.format(args.start_epoch)) best_map = [0] for epoch in range(args.start_epoch, args.epochs): mix_ratio = 1.0 if args.mixup: # TODO(zhreshold) only support evenly mixup now, target generator needs to be modified otherwise train_data._dataset.set_mixup(np.random.uniform, 0.5, 0.5) mix_ratio = 0.5 if epoch >= args.epochs - args.no_mixup_epochs: train_data._dataset.set_mixup(None) mix_ratio = 1.0 while lr_steps and epoch >= lr_steps[0]: new_lr = trainer.learning_rate * lr_decay lr_steps.pop(0) trainer.set_learning_rate(new_lr) logger.info("[Epoch {}] Set learning rate to {}".format(epoch, new_lr)) for metric in metrics: metric.reset() tic = time.time() btic = time.time() net.hybridize(static_alloc=True) base_lr = trainer.learning_rate for i, batch in enumerate(train_data): if epoch == 0 and i <= lr_warmup: # adjust based on real percentage new_lr = base_lr * get_lr_at_iter(i / lr_warmup) if new_lr != trainer.learning_rate: if i % args.log_interval == 0: logger.info('[Epoch 0 Iteration {}] Set learning rate to {}'.format(i, new_lr)) trainer.set_learning_rate(new_lr) batch = split_and_load(batch, ctx_list=ctx) batch_size = len(batch[0]) losses = [] metric_losses = [[] for _ in metrics] add_losses = [[] for _ in metrics2] with autograd.record(): for data, label, rpn_cls_targets, rpn_box_targets, rpn_box_masks in zip(*batch): gt_label = label[:, :, 4:5] gt_box = label[:, :, :4] cls_pred, box_pred, roi, samples, matches, rpn_score, rpn_box, anchors = net(data, gt_box) # losses of rpn rpn_score = rpn_score.squeeze(axis=-1) num_rpn_pos = (rpn_cls_targets >= 0).sum() rpn_loss1 = rpn_cls_loss(rpn_score, rpn_cls_targets, rpn_cls_targets >= 0) * rpn_cls_targets.size / num_rpn_pos rpn_loss2 = rpn_box_loss(rpn_box, rpn_box_targets, rpn_box_masks) * rpn_box.size / num_rpn_pos # rpn overall loss, use sum rather than average rpn_loss = rpn_loss1 + rpn_loss2 # generate targets for rcnn cls_targets, box_targets, box_masks = net.target_generator(roi, samples, matches, gt_label, gt_box) # losses of rcnn num_rcnn_pos = (cls_targets >= 0).sum() rcnn_loss1 = rcnn_cls_loss(cls_pred, cls_targets, cls_targets >= 0) * cls_targets.size / cls_targets.shape[0] / num_rcnn_pos rcnn_loss2 = rcnn_box_loss(box_pred, box_targets, box_masks) * box_pred.size / box_pred.shape[0] / num_rcnn_pos rcnn_loss = rcnn_loss1 + rcnn_loss2 # overall losses losses.append(rpn_loss.sum() * mix_ratio + rcnn_loss.sum() * mix_ratio) metric_losses[0].append(rpn_loss1.sum() * mix_ratio) metric_losses[1].append(rpn_loss2.sum() * mix_ratio) metric_losses[2].append(rcnn_loss1.sum() * mix_ratio) metric_losses[3].append(rcnn_loss2.sum() * mix_ratio) add_losses[0].append([[rpn_cls_targets, rpn_cls_targets>=0], [rpn_score]]) add_losses[1].append([[rpn_box_targets, rpn_box_masks], [rpn_box]]) add_losses[2].append([[cls_targets], [cls_pred]]) add_losses[3].append([[box_targets, box_masks], [box_pred]]) autograd.backward(losses) for metric, record in zip(metrics, metric_losses): metric.update(0, record) for metric, records in zip(metrics2, add_losses): for pred in records: metric.update(pred[0], pred[1]) trainer.step(batch_size) # update metrics if args.log_interval and not (i + 1) % args.log_interval: # msg = ','.join(['{}={:.3f}'.format(*metric.get()) for metric in metrics]) msg = ','.join(['{}={:.3f}'.format(*metric.get()) for metric in metrics + metrics2]) logger.info('[Epoch {}][Batch {}], Speed: {:.3f} samples/sec, {}'.format( epoch, i, args.log_interval * batch_size/(time.time()-btic), msg)) btic = time.time() msg = ','.join(['{}={:.3f}'.format(*metric.get()) for metric in metrics]) logger.info('[Epoch {}] Training cost: {:.3f}, {}'.format( epoch, (time.time()-tic), msg)) # if not (epoch + 1) % args.val_interval: # # consider reduce the frequency of validation to save time # map_name, mean_ap = validate(net, val_data, ctx, eval_metric) # val_msg = '\n'.join(['{}={}'.format(k, v) for k, v in zip(map_name, mean_ap)]) # logger.info('[Epoch {}] Validation: \n{}'.format(epoch, val_msg)) # current_map = float(mean_ap[-1]) # else: # current_map = 0. current_map = 0 save_params(net, logger, best_map, current_map, epoch, args.save_interval, args.save_prefix) if __name__ == '__main__': args = parse_args() # fix seed for mxnet, numpy and python builtin random generator. gutils.random.seed(args.seed) # training contexts ctx = [mx.gpu(int(i)) for i in args.gpus.split(',') if i.strip()] ctx = ctx if ctx else [mx.cpu()] args.batch_size = len(ctx) # 1 batch per device # network net_name = '_'.join(('faster_rcnn', args.network, args.dataset)) args.save_prefix += net_name net = get_model(net_name, pretrained_base=True) if args.resume.strip(): net.load_parameters(args.resume.strip()) else: for param in net.collect_params().values(): if param._data is not None: continue param.initialize() net.collect_params().reset_ctx(ctx) # training data train_dataset, val_dataset, eval_metric = get_dataset(args.dataset, args) train_data, val_data = get_dataloader( net, train_dataset, val_dataset, args.batch_size, args.num_workers) # training train(net, train_data, val_data, eval_metric, ctx, args)
View Code

 

檢測部分,是在demo 下修改的,填了幾個引數,可以用lst檔案遍歷了,用cv2畫圖,不用那個matplotlib了

"""Faster RCNN Demo script."""
import os
import argparse
import mxnet as mx
import gluoncv as gcv
from gluoncv.data.transforms import presets
from matplotlib import pyplot as plt
import cv2

font = cv2.FONT_HERSHEY_SIMPLEX

def parse_args():
    parser = argparse.ArgumentParser(description='Test with Faster RCNN networks.')
    parser.add_argument('--network', type=str, default='faster_rcnn_resnet50_v1b_coco',
                        help="Faster RCNN full network name")
    parser.add_argument('--images', type=str, default='',
                        help='Test images, use comma to split multiple.')
    parser.add_argument('--gpus', type=str, default='',
                        help='Training with GPUs, you can specify 1,3 for example.')
    parser.add_argument('--pretrained', type=str, default='True',
                        help='Load weights from previously saved parameters. You can specify parameter file name.')
    parser.add_argument('--thresh', type=float, default=0.5,
                        help='Threshold of object score when visualize the bboxes.')
    # add_lst
    parser.add_argument('--lst', type=str,default='',help="predict's lst file")
    args = parser.parse_args()
    return args

if __name__ == '__main__':
    args = parse_args()
    # context list
    ctx = [mx.gpu(int(i)) for i in args.gpus.split(',') if i.strip()]
    ctx = [mx.cpu()] if not ctx else ctx

    # grab some image if not specified
    if not args.images.strip() and args.lst=='':
        gcv.utils.download('https://github.com/dmlc/web-data/blob/master/' +
                           'gluoncv/detection/biking.jpg?raw=true', 'biking.jpg')
        image_list = ['biking.jpg']
    else:
        image_list = [x.strip() for x in args.images.split(',') if x.strip()]
    
    cnt = 0
    if args.lst!='':
        print(args.lst)
        file = open('val_front_0913.lst')
        image_list = []
        for line in file:
            line = line.split('\t')
            print('/mnt/hdfs-data-4/data/jian.yin/val_front_0913/'+line[-1][:-1])
            image_list.append('/mnt/hdfs-data-4/data/jian.yin/val_front_0913/'+line[-1][:-1])
            cnt+=1
        print 'sum of pic ',cnt
    
    if args.pretrained.lower() in ['true', '1', 'yes', 't']:
        net = gcv.model_zoo.get_model(args.network, pretrained=True)
    else:
        net = gcv.model_zoo.get_model(args.network, pretrained=False, pretrained_base=False)
        net.load_parameters(args.pretrained)
    net.set_nms(0.3, 200)
    net.collect_params().reset_ctx(ctx = ctx)

    ax = None
    
    # write plt.txt
    fw = open('draw/plt.txt','w')
    dict = {}
    cnt1 = 0
    for image in image_list:
        dict['url'] = image
        bbox_list = []
        x, img = presets.rcnn.load_test(image, short=net.short, max_size=net.max_size)
        img_h = img.shape[0]
        img_w = img.shape[1]
        
        x = x.as_in_context(ctx[0])
        ids, scores, bboxes = [xx[0].asnumpy() for xx in net(x)]

        original_img = cv2.imread(image)
        original_img_h = original_img.shape[0]
        original_img_w = original_img.shape[1]
        
        for i in range(scores.shape[0]):
            if scores[i] > args.thresh:
                x1 = int(bboxes[i][0]*original_img_h/img_h)
                y1 = int(bboxes[i][1]*original_img_w/img_w)
                x2 = int(bboxes[i][2]*original_img_h/img_h)
                y2 = int(bboxes[i][3]*original_img_w/img_w)
                
                bbox_list.append((float(scores[i]),x1,y1,x2,y2))
        dict['bbox'] = bbox_list
        fw.write(str(dict)+'\n')
        cnt1+=1
        print 'The last ',cnt-cnt1
    fw.close()
#                 cv2.rectangle(original_img, (x1, y1), (x2, y2), (255,0,0), 3)
#                 cv2.putText(original_img,'person '+str(scores[i]),(x1,y1),font,0.5,(255,0,0),2)
#                 cv2.imwrite('draw/'+str(cnt)+'.jpg', original_img)
                
        
        
#         print(bboxes)
#         ax = gcv.utils.viz.plot_bbox(img, bboxes, scores, ids, thresh=args.thresh,
#                                      class_names=net.classes, ax=ax)
#         plt.savefig(str(cnt)+'predict.jpg')
#         cnt+=1
#         plt.show()

把得分情況,錨框位置都寫在檔案裡了,不用每次跑模型來得到,想怎麼都可以了。plt.py

import cv2
import os
font = cv2.FONT_HERSHEY_SIMPLEX

file = open('plt.txt')
cnt = 1
for line in file:
    dict = eval(line)
    url = dict['url']
    bbox = dict['bbox']
    img = cv2.imread(url)
    for i in range(len(bbox)):
        score = bbox[i][0]
        score = '%.2f' % score
        x1 = bbox[i][1]
        y1 = bbox[i][2]
        x2 = bbox[i][3]
        y2 = bbox[i][4]
        cv2.rectangle(img, (x1, y1), (x2, y2), (255,0,0), 3)
        cv2.putText(img,'person '+str(score),(x1,y1),font,0.5,(255,0,0),2)
    url = url.split('/')
    x_url = url[5]+'/'+url[6]+'/'+url[7]+'/'+url[8]
    if not os.path.exists(url[5]+'/'+url[6]+'/'+url[7]+'/'):
        os.makedirs(url[5]+'/'+url[6]+'/'+url[7]+'/')
    cv2.imwrite(x_url, img)
    print('The last ',6137-cnt)
    cnt+=1