zoukankan      html  css  js  c++  java
  • 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}:	{:.4f}
    '.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 = '
    '.join(['{}={}'.format(k, v) for k, v in zip(map_name, mean_ap)])
                
                
                
    #             logger.info('[Epoch {}] Validation: 
    {}'.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('	')
                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)+'
    ')
            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
  • 相关阅读:
    python2和python3中TestSuite().addTest的区别
    python2和python3中range的区别
    WebDriverAgent安装
    Charles IOS https抓包
    Python和 pytest的异常处理
    Python中yaml和json文件的读取和应用
    Python 获取当前文件所在路径
    自建包函数调用
    python的logging,将log保存到文件
    泛型
  • 原文地址:https://www.cnblogs.com/TreeDream/p/10180651.html
Copyright © 2011-2022 走看看