コード例 #1
0
    def main():
        parser = argparse.ArgumentParser()
        parser.add_argument('-s', '--dataset', type=str, choices=DatasetBase.OPTIONS, required=True, help='name of dataset')
        parser.add_argument('-b', '--backbone', type=str, choices=BackboneBase.OPTIONS, required=True, help='name of backbone model')
        parser.add_argument('-c', '--checkpoint', type=str, required=True, help='path to checkpoint')
        parser.add_argument('-p', '--probability_threshold', type=float, default=0.6, help='threshold of detection probability')
        parser.add_argument('--image_min_side', type=float, help='default: {:g}'.format(Config.IMAGE_MIN_SIDE))
        parser.add_argument('--image_max_side', type=float, help='default: {:g}'.format(Config.IMAGE_MAX_SIDE))
        parser.add_argument('--anchor_ratios', type=str, help='default: "{!s}"'.format(Config.ANCHOR_RATIOS))
        parser.add_argument('--anchor_sizes', type=str, help='default: "{!s}"'.format(Config.ANCHOR_SIZES))
        parser.add_argument('--pooler_mode', type=str, choices=Pooler.OPTIONS, help='default: {.value:s}'.format(Config.POOLER_MODE))
        parser.add_argument('--rpn_pre_nms_top_n', type=int, help='default: {:d}'.format(Config.RPN_PRE_NMS_TOP_N))
        parser.add_argument('--rpn_post_nms_top_n', type=int, help='default: {:d}'.format(Config.RPN_POST_NMS_TOP_N))
        args = parser.parse_args()

        dataset_name = args.dataset
        backbone_name = args.backbone
        path_to_checkpoint = args.checkpoint
        prob_thresh = args.probability_threshold

        Config.setup(image_min_side=args.image_min_side, image_max_side=args.image_max_side,
                     anchor_ratios=args.anchor_ratios, anchor_sizes=args.anchor_sizes, pooler_mode=args.pooler_mode,
                     rpn_pre_nms_top_n=args.rpn_pre_nms_top_n, rpn_post_nms_top_n=args.rpn_post_nms_top_n)

        print('Arguments:')
        for k, v in vars(args).items():
            print(f'\t{k} = {v}')
        print(Config.describe())

        _infer_websocket(path_to_checkpoint, dataset_name, backbone_name, prob_thresh)
コード例 #2
0
    def main():
        parser = argparse.ArgumentParser()
        parser.add_argument('input', type=str, help='path to input image')
        parser.add_argument('output', type=str, help='path to output result image')
        parser.add_argument('-c', '--checkpoint', type=str, required=True, help='path to checkpoint')
        parser.add_argument('-s', '--dataset', type=str, choices=DatasetBase.OPTIONS, required=True, help='name of dataset')
        parser.add_argument('-b', '--backbone', type=str, choices=BackboneBase.OPTIONS, required=True, help='name of backbone model')
        parser.add_argument('-p', '--probability_threshold', type=float, default=0.6, help='threshold of detection probability')
        parser.add_argument('--image_min_side', type=float, help='default: {:g}'.format(Config.IMAGE_MIN_SIDE))
        parser.add_argument('--image_max_side', type=float, help='default: {:g}'.format(Config.IMAGE_MAX_SIDE))
        parser.add_argument('--anchor_ratios', type=str, help='default: "{!s}"'.format(Config.ANCHOR_RATIOS))
        parser.add_argument('--anchor_scales', type=str, help='default: "{!s}"'.format(Config.ANCHOR_SCALES))
        parser.add_argument('--pooling_mode', type=str, choices=ROIWrapper.OPTIONS, help='default: {.value:s}'.format(Config.POOLING_MODE))
        parser.add_argument('--rpn_pre_nms_top_n', type=int, help='default: {:d}'.format(Config.RPN_PRE_NMS_TOP_N))
        parser.add_argument('--rpn_post_nms_top_n', type=int, help='default: {:d}'.format(Config.RPN_POST_NMS_TOP_N))
        args = parser.parse_args()

        path_to_input_image = args.input
        path_to_output_image = args.output
        path_to_checkpoint = args.checkpoint
        dataset_name = args.dataset
        backbone_name = args.backbone
        prob_thresh = args.probability_threshold

        os.makedirs(os.path.join(os.path.curdir, os.path.dirname(path_to_output_image)), exist_ok=True)

        Config.setup(image_min_side=args.image_min_side, image_max_side=args.image_max_side,
                     anchor_ratios=args.anchor_ratios, anchor_scales=args.anchor_scales, pooling_mode=args.pooling_mode,
                     rpn_pre_nms_top_n=args.rpn_pre_nms_top_n, rpn_post_nms_top_n=args.rpn_post_nms_top_n)

        print('Arguments:')
        for k, v in vars(args).items():
            print(f'\t{k} = {v}')
        print(Config.describe())

        _infer(path_to_input_image, path_to_output_image, path_to_checkpoint, dataset_name, backbone_name, prob_thresh)
コード例 #3
0
ファイル: eval.py プロジェクト: wyh20000305/easy-fpn.pytorch
    def main():
        parser = argparse.ArgumentParser()
        parser.add_argument('checkpoint', type=str, help='path to evaluating checkpoint')
        parser.add_argument('-s', '--dataset', type=str, choices=DatasetBase.OPTIONS, required=True, help='name of dataset')
        parser.add_argument('-b', '--backbone', type=str, choices=BackboneBase.OPTIONS, required=True, help='name of backbone model')
        parser.add_argument('-d', '--data_dir', type=str, default='./data', help='path to data directory')
        parser.add_argument('--image_min_side', type=float, help='default: {:g}'.format(Config.IMAGE_MIN_SIDE))
        parser.add_argument('--image_max_side', type=float, help='default: {:g}'.format(Config.IMAGE_MAX_SIDE))
        parser.add_argument('--anchor_ratios', type=str, help='default: "{!s}"'.format(Config.ANCHOR_RATIOS))
        parser.add_argument('--anchor_scales', type=str, help='default: "{!s}"'.format(Config.ANCHOR_SCALES))
        parser.add_argument('--pooling_mode', type=str, choices=ROIWrapper.OPTIONS, help='default: {.value:s}'.format(Config.POOLING_MODE))
        parser.add_argument('--rpn_pre_nms_top_n', type=int, help='default: {:d}'.format(Config.RPN_PRE_NMS_TOP_N))
        parser.add_argument('--rpn_post_nms_top_n', type=int, help='default: {:d}'.format(Config.RPN_POST_NMS_TOP_N))
        args = parser.parse_args()

        path_to_checkpoint = args.checkpoint
        dataset_name = args.dataset
        backbone_name = args.backbone
        path_to_data_dir = args.data_dir

        path_to_results_dir = os.path.join(os.path.dirname(path_to_checkpoint), 'results-{:s}-{:s}-{:s}'.format(
            time.strftime('%Y%m%d%H%M%S'), path_to_checkpoint.split(os.path.sep)[-1].split(os.path.curdir)[0],
            str(uuid.uuid4()).split('-')[0]))
        os.makedirs(path_to_results_dir)

        Config.setup(image_min_side=args.image_min_side, image_max_side=args.image_max_side,
                     anchor_ratios=args.anchor_ratios, anchor_scales=args.anchor_scales, pooling_mode=args.pooling_mode,
                     rpn_pre_nms_top_n=args.rpn_pre_nms_top_n, rpn_post_nms_top_n=args.rpn_post_nms_top_n)

        Log.initialize(os.path.join(path_to_results_dir, 'eval.log'))
        Log.i('Arguments:')
        for k, v in vars(args).items():
            Log.i(f'\t{k} = {v}')
        Log.i(Config.describe())

        _eval(path_to_checkpoint, dataset_name, backbone_name, path_to_data_dir, path_to_results_dir)
コード例 #4
0
    def main():
        parser = argparse.ArgumentParser()
        parser.add_argument('--image_min_side',
                            type=float,
                            help='default: {:g}'.format(Config.IMAGE_MIN_SIDE))
        parser.add_argument('--image_max_side',
                            type=float,
                            help='default: {:g}'.format(Config.IMAGE_MAX_SIDE))
        parser.add_argument('--anchor_ratios',
                            type=str,
                            help='default: "{!s}"'.format(
                                Config.ANCHOR_RATIOS))
        parser.add_argument('--anchor_sizes',
                            type=str,
                            help='default: "{!s}"'.format(Config.ANCHOR_SIZES))
        parser.add_argument('--pooler_mode',
                            type=str,
                            choices=Pooler.OPTIONS,
                            help='default: {.value:s}'.format(
                                Config.POOLER_MODE))
        parser.add_argument('--rpn_pre_nms_top_n',
                            type=int,
                            help='default: {:d}'.format(
                                Config.RPN_PRE_NMS_TOP_N))
        parser.add_argument('--rpn_post_nms_top_n',
                            type=int,
                            help='default: {:d}'.format(
                                Config.RPN_POST_NMS_TOP_N))

        args = parser.parse_args()

        input_root = '/home/mmlab/CCTV_Server/models/detectors/FasterRCNN/frames'
        output_root = input_root + '_output'
        path_to_checkpoint = '/home/mmlab/CCTV_Server/models/detectors/FasterRCNN/checkpoints/obstacle/model-90000.pth'
        dataset_name = 'obstacle'
        backbone_name = 'resnet101'
        prob_thresh = 0.6
        Config.setup(image_min_side=args.image_min_side,
                     image_max_side=args.image_max_side,
                     anchor_ratios=args.anchor_ratios,
                     anchor_sizes=args.anchor_sizes,
                     pooler_mode=args.pooler_mode,
                     rpn_pre_nms_top_n=args.rpn_pre_nms_top_n,
                     rpn_post_nms_top_n=args.rpn_post_nms_top_n)

        print('Arguments:')
        for k, v in vars(args).items():
            print(f'\t{k} = {v}')
        print(Config.describe())

        os.makedirs(output_root, exist_ok=True)

        input_sub_dirnames = [
            directory for directory in os.listdir(input_root)
            if os.path.isdir(os.path.join(input_root, directory))
        ]
        dataset_class = DatasetBase.from_name(dataset_name)
        backbone = BackboneBase.from_name(backbone_name)(pretrained=False)
        model = Model(backbone,
                      dataset_class.num_classes(),
                      pooler_mode=Config.POOLER_MODE,
                      anchor_ratios=Config.ANCHOR_RATIOS,
                      anchor_sizes=Config.ANCHOR_SIZES,
                      rpn_pre_nms_top_n=Config.RPN_PRE_NMS_TOP_N,
                      rpn_post_nms_top_n=Config.RPN_POST_NMS_TOP_N).cuda()
        model.load(path_to_checkpoint)

        for sub_dir in input_sub_dirnames:
            input_sub_dirpath = os.path.join(input_root, sub_dir)
            output_sub_dirpath = os.path.join(output_root, sub_dir)

            filenames = [
                image_basename(f) for f in os.listdir(input_sub_dirpath)
                if is_image(f)
            ]
            for filename in filenames:
                path_to_input_image = image_path(input_sub_dirpath, filename,
                                                 '.jpg')
                # path_to_input_image = '/faster-RCNN/frames/1_360p/1_360p_0001.jpg'
                path_to_output_image = image_path(output_sub_dirpath, filename,
                                                  '.jpg')
                # path_to_output_image = '/faster-RCNN/frames_output/1_360p/1_360p_0001.jpg'

                os.makedirs(os.path.join(
                    os.path.curdir, os.path.dirname(path_to_output_image)),
                            exist_ok=True)

                with torch.no_grad():
                    image = transforms.Image.open(path_to_input_image)
                    image_tensor, scale = dataset_class.preprocess(
                        image, Config.IMAGE_MIN_SIDE, Config.IMAGE_MAX_SIDE)

                    detection_bboxes, detection_classes, detection_probs, _, _ = \
                        model.eval().forward(image_tensor.unsqueeze(dim=0).cuda())
                    detection_bboxes /= scale

                    kept_indices = detection_probs > prob_thresh
                    detection_bboxes = detection_bboxes[kept_indices]
                    detection_classes = detection_classes[kept_indices]
                    detection_probs = detection_probs[kept_indices]

                    draw = ImageDraw.Draw(image)

                    for bbox, cls, prob in zip(detection_bboxes.tolist(),
                                               detection_classes.tolist(),
                                               detection_probs.tolist()):
                        color = random.choice([
                            'red', 'green', 'blue', 'yellow', 'purple', 'white'
                        ])
                        bbox = BBox(left=bbox[0],
                                    top=bbox[1],
                                    right=bbox[2],
                                    bottom=bbox[3])
                        category = dataset_class.LABEL_TO_CATEGORY_DICT[cls]

                        draw.rectangle(
                            ((bbox.left, bbox.top), (bbox.right, bbox.bottom)),
                            outline=color)
                        draw.text((bbox.left, bbox.top),
                                  text=f'{category:s} {prob:.3f}',
                                  fill=color)

                    image.save(path_to_output_image)
                    print(f'Output image is saved to {path_to_output_image}')
コード例 #5
0
    def main():
        parser = argparse.ArgumentParser()
        '''
        python infer.py -s=coco2017 -b=resnet101 -c=model-180000.pth --image_min_side=800 --image_max_side=1333 --anchor_sizes="[64, 128, 256, 512]" --rpn_post_nms_top_n=1000 field.jpg out_field.jpg
        
        这个model-180000.pth对应的训练数据是coco2017, 因此不能使用voc2007.
        '''
        #指定映射参数 > COOC模型以及验证
        # parser.add_argument('-s', '--dataset', type=str, default='coco2017',choices=DatasetBase.OPTIONS, help='name of dataset')
        # parser.add_argument('-b', '--backbone', type=str, default='resnet101',choices=BackboneBase.OPTIONS,  help='name of backbone model')
        # parser.add_argument('-c', '--checkpoint', type=str, default='model-180000.pth', help='path to checkpoint')

        #>VOC2007
        parser.add_argument('-s',
                            '--dataset',
                            type=str,
                            default='voc2007',
                            choices=DatasetBase.OPTIONS,
                            help='name of dataset')
        parser.add_argument('-b',
                            '--backbone',
                            type=str,
                            default='resnet101',
                            choices=BackboneBase.OPTIONS,
                            help='name of backbone model')
        model_path = 'outputs/checkpoints-20190702190950-voc2007-resnet101-f8049269/model-22500.pth'
        parser.add_argument('-c',
                            '--checkpoint',
                            type=str,
                            default=model_path,
                            help='path to checkpoint')

        parser.add_argument('-p',
                            '--probability_threshold',
                            type=float,
                            default=0.6,
                            help='threshold of detection probability')
        parser.add_argument('--image_min_side',
                            default=600,
                            type=float,
                            help='default: {:g}'.format(Config.IMAGE_MIN_SIDE))
        parser.add_argument('--image_max_side',
                            default=1000,
                            type=float,
                            help='default: {:g}'.format(Config.IMAGE_MAX_SIDE))
        parser.add_argument('--anchor_ratios',
                            type=str,
                            help='default: "{!s}"'.format(
                                Config.ANCHOR_RATIOS))
        parser.add_argument('--anchor_sizes',
                            type=str,
                            default="[128, 256, 512]",
                            help='default: "{!s}"'.format(Config.ANCHOR_SIZES))
        parser.add_argument('--pooler_mode',
                            type=str,
                            choices=Pooler.OPTIONS,
                            help='default: {.value:s}'.format(
                                Config.POOLER_MODE))
        parser.add_argument('--rpn_pre_nms_top_n',
                            type=int,
                            help='default: {:d}'.format(
                                Config.RPN_PRE_NMS_TOP_N))
        parser.add_argument('--rpn_post_nms_top_n',
                            type=int,
                            default=1000,
                            help='default: {:d}'.format(
                                Config.RPN_POST_NMS_TOP_N))
        #最后两个是位置参数,没有指定映射参数
        #parser.add_argument('input', type=str, default='field.jpg',help='path to input image')
        #parser.add_argument('output', type=str, default='out_filed.jpg',help='path to output result image')
        args = parser.parse_args()

        path_to_input_image = '1.jpg'
        path_to_output_image = 'out_coco.jpg'
        dataset_name = args.dataset
        backbone_name = args.backbone
        path_to_checkpoint = args.checkpoint
        prob_thresh = args.probability_threshold

        os.makedirs(os.path.join(os.path.curdir,
                                 os.path.dirname(path_to_output_image)),
                    exist_ok=True)

        #Config是由类变量和类方法构成的属性配置器
        Config.setup(image_min_side=args.image_min_side,
                     image_max_side=args.image_max_side,
                     anchor_ratios=args.anchor_ratios,
                     anchor_sizes=args.anchor_sizes,
                     pooler_mode=args.pooler_mode,
                     rpn_pre_nms_top_n=args.rpn_pre_nms_top_n,
                     rpn_post_nms_top_n=args.rpn_post_nms_top_n)

        print('Arguments:')
        for k, v in vars(args).items():
            print(f'\t{k} = {v}')
        print(Config.describe())

        _infer(path_to_input_image, path_to_output_image, path_to_checkpoint,
               dataset_name, backbone_name, prob_thresh)