基于PaddleDetection目标检测labelme标注自动获取

在百度的PaddleDetection项目的基础上实现目标检测labelme标注的自动获取，需要先训练一个模型，然后通过这个模型去标注，最后用labelme进行微调

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import os
import sys
import json
import io
import base64

# add python path of PadleDetection to sys.path
from ppdet.data.source.category import get_categories
from ppdet.optimizer import ModelEMA
from ppdet.utils.checkpoint import load_pretrain_weight
from ppdet.utils.visualizer import visualize_results

parent_path = os.path.abspath(os.path.join(__file__, *(['..'] * 2)))
sys.path.insert(0, parent_path)

import warnings
warnings.filterwarnings('ignore')
import glob
import paddle
from ppdet.core.workspace import load_config, merge_config, create
from ppdet.utils.check import check_gpu, check_npu, check_version, check_config
from ppdet.utils.cli import ArgsParser
from ppdet.slim import build_slim_model
from ppdet.metrics import  get_infer_results
from PIL import Image, ImageOps

import numpy as np
from ppdet.utils.logger import setup_logger
logger = setup_logger('train')


def parse_args():
    parser = ArgsParser()
    parser.add_argument(
        "--infer_dir",
        type=str,
        default=r"C:\Users\86187\Desktop\classifyUnrecognizedESAll",
        help="Directory for images to perform inference on.")
    parser.add_argument(
        "--infer_img",
        type=str,
        default=None,
        help="Image path, has higher priority over --infer_dir")
    parser.add_argument(
        "--output_dir",
        type=str,
        default="python_infer_output",
        help="Directory for storing the output visualization files.")
    parser.add_argument(
        "--draw_threshold",
        type=float,
        default=0.5,
        help="Threshold to reserve the result for visualization.")
    parser.add_argument(
        "--slim_config",
        default=None,
        type=str,
        help="Configuration file of slim method.")
    parser.add_argument(
        "--use_vdl",
        type=bool,
        default=False,
        help="Whether to record the data to VisualDL.")
    parser.add_argument(
        '--vdl_log_dir',
        type=str,
        default="vdl_log_dir/image",
        help='VisualDL logging directory for image.')
    parser.add_argument(
        "--save_txt",
        type=bool,
        default=False,
        help="Whether to save inference result in txt.")
    args = parser.parse_args()
    return args


def img_arr_to_b64(img_pil):
    # img_pil = Image.fromarray(img_arr)
    f = io.BytesIO()
    img_pil.save(f, format="PNG")
    img_bin = f.getvalue()
    if hasattr(base64, "encodebytes"):
        img_b64 = base64.encodebytes(img_bin)
    else:
        img_b64 = base64.encodestring(img_bin)
    return img_b64


def _get_save_image_name(output_dir, image_path):
    """
    Get save image name from source image path.
    """
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
    image_name = os.path.split(image_path)[-1]
    name, ext = os.path.splitext(image_name)
    return os.path.join(output_dir, "{}".format(name)) + ext


def get_test_images(infer_dir, infer_img):
    """
    Get image path list in TEST mode
    """
    assert infer_img is not None or infer_dir is not None, \
        "--infer_img or --infer_dir should be set"
    assert infer_img is None or os.path.isfile(infer_img), \
            "{} is not a file".format(infer_img)
    assert infer_dir is None or os.path.isdir(infer_dir), \
            "{} is not a directory".format(infer_dir)

    # infer_img has a higher priority
    if infer_img and os.path.isfile(infer_img):
        return [infer_img]

    images = set()
    infer_dir = os.path.abspath(infer_dir)
    assert os.path.isdir(infer_dir), \
        "infer_dir {} is not a directory".format(infer_dir)
    exts = ['jpg', 'jpeg', 'png', 'bmp']
    exts += [ext.upper() for ext in exts]
    for ext in exts:
        images.update(glob.glob('{}/*.{}'.format(infer_dir, ext)))
    images = list(images)

    assert len(images) > 0, "no image found in {}".format(infer_dir)
    logger.info("Found {} inference images in total.".format(len(images)))

    return images


def _coco17_category():
    """
    Get class id to category id map and category id
    to category name map of COCO2017 dataset

    """
    clsid2catid = {
        1: 1,
        2: 2,
        3: 3,
        4: 4,
        5: 5,
        6: 6,
        7: 7,
        8: 8,
        9: 9,
        10: 10,
        11: 11,
        12: 13,
        13: 14,
        14: 15,
        15: 16,
        16: 17,
        17: 18,
        18: 19,
        19: 20,
        20: 21,
        21: 22,
        22: 23,
        23: 24,
        24: 25,
        25: 27,
        26: 28,
        27: 31,
        28: 32,
        29: 33,
        30: 34,
        31: 35,
        32: 36,
        33: 37,
        34: 38,
        35: 39,
        36: 40,
        37: 41,
        38: 42,
        39: 43,
        40: 44,
        41: 46,
        42: 47,
        43: 48,
        44: 49,
        45: 50,
        46: 51,
        47: 52,
        48: 53,
        49: 54,
        50: 55,
        51: 56,
        52: 57,
        53: 58,
        54: 59,
        55: 60,
        56: 61,
        57: 62,
        58: 63,
        59: 64,
        60: 65,
        61: 67,
        62: 70,
        63: 72,
        64: 73,
        65: 74,
        66: 75,
        67: 76,
        68: 77,
        69: 78,
        70: 79,
        71: 80,
        72: 81,
        73: 82,
        74: 84,
        75: 85,
        76: 86,
        77: 87,
        78: 88,
        79: 89,
        80: 90
    }

    catid2name = {
        0: 'background',
        1: 'barCode',
        2: '5',
        3: '4',
        4: '8',
        5: '6',
        6: '9',
        7: '3',
        8: '0',
        9: '2',
        10: '1',
        11: 'fire hydrant',
        13: 'stop sign',
        14: 'parking meter',
        15: 'bench',
        16: 'bird',
        17: 'cat',
        18: 'dog',
        19: 'horse',
        20: 'sheep',
        21: 'cow',
        22: 'elephant',
        23: 'bear',
        24: 'zebra',
        25: 'giraffe',
        27: 'backpack',
        28: 'umbrella',
        31: 'handbag',
        32: 'tie',
        33: 'suitcase',
        34: 'frisbee',
        35: 'skis',
        36: 'snowboard',
        37: 'sports ball',
        38: 'kite',
        39: 'baseball bat',
        40: 'baseball glove',
        41: 'skateboard',
        42: 'surfboard',
        43: 'tennis racket',
        44: 'bottle',
        46: 'wine glass',
        47: 'cup',
        48: 'fork',
        49: 'knife',
        50: 'spoon',
        51: 'bowl',
        52: 'banana',
        53: 'apple',
        54: 'sandwich',
        55: 'orange',
        56: 'broccoli',
        57: 'carrot',
        58: 'hot dog',
        59: 'pizza',
        60: 'donut',
        61: 'cake',
        62: 'chair',
        63: 'couch',
        64: 'potted plant',
        65: 'bed',
        67: 'dining table',
        70: 'toilet',
        72: 'tv',
        73: 'laptop',
        74: 'mouse',
        75: 'remote',
        76: 'keyboard',
        77: 'cell phone',
        78: 'microwave',
        79: 'oven',
        80: 'toaster',
        81: 'sink',
        82: 'refrigerator',
        84: 'book',
        85: 'clock',
        86: 'vase',
        87: 'scissors',
        88: 'teddy bear',
        89: 'hair drier',
        90: 'toothbrush'
    }

    clsid2catid = {k - 1: v for k, v in clsid2catid.items()}
    catid2name.pop(0)

    return clsid2catid, catid2name


def getImagesLabels(image, bboxes, im_id, save_image_name, catid2name, draw_threshold):

    name = save_image_name.split(".")
    save_json_name = name[0] + '.json'

    w, h = image.size

    # label的写入
    dst_data = dict()
    dst_data['version'] = "4.5.9"
    dst_data['flags'] = {}
    dst_data['shapes'] = []
    dst_data['imagePath'] = save_image_name.split("\\")[-1]
    dst_data['imageData'] = img_arr_to_b64(image).decode('utf-8')
    dst_data['imageHeight'] = h
    dst_data['imageWidth'] = w

    for dt in np.array(bboxes):
        if im_id != dt['image_id']:
            continue

        catid, bbox, score = dt['category_id'], dt['bbox'], dt['score']
        if score < draw_threshold:
            continue

        xmin, ymin, w, h = bbox
        label = catid2name[catid]

        shape = dict()
        shape['label'] = label
        shape['points'] = [[xmin, ymin], [xmin + w, ymin + h]]
        shape['group_id'] = None
        shape['shape_type'] = "rectangle"
        shape['flags'] = {}
        dst_data['shapes'].append(shape)
    print(save_json_name)
    json.dump(dst_data, open(save_json_name, 'w'), indent=4)


def run(FLAGS, cfg):
    draw_threshold = 0.5
    output_dir = 'output'
    labels_dir = 'labelme'
    # build data loader
    mode = 'test'
    dataset = cfg['{}Dataset'.format(mode.capitalize())]

    # get inference images
    images = get_test_images(FLAGS.infer_dir, FLAGS.infer_img)

    # build model
    model = create(cfg.architecture)

    # normalize params for deploy
    model.load_meanstd(cfg['TestReader']['sample_transforms'])

    use_ema = ('use_ema' in cfg and cfg['use_ema'])
    if use_ema:
        ema_decay = cfg.get('ema_decay', 0.9998)
        cycle_epoch = cfg.get('cycle_epoch', -1)
        ema = ModelEMA(
            model,
            decay=ema_decay,
            use_thres_step=True,
            cycle_epoch=cycle_epoch)

    # load weights
    load_pretrain_weight(model, cfg.weights)

    # # predict
    dataset.set_images(images)
    loader = create('TestReader')(dataset, 0)

    imid2path = dataset.get_imid2path()

    anno_file = dataset.get_anno()
    clsid2catid, catid2name = _coco17_category()
    # print(clsid2catid)
    # print(catid2name)

    # Run Infer
    model.eval()
    status = {'mode': 'test'}
    results = []
    for step_id, data in enumerate(loader):
        status['step_id'] = step_id
        # forward
        outs = model(data)

        # print(outs)

        for key in ['im_shape', 'scale_factor', 'im_id']:
            outs[key] = data[key]
        for key, value in outs.items():
            if hasattr(value, 'numpy'):  # hasattr() 函数用于判断对象是否包含对应的属性。hasattr(object, name)
                outs[key] = value.numpy()
        results.append(outs)
        # sniper

    for outs in results:
        batch_res = get_infer_results(outs, clsid2catid)
        bbox_num = outs['bbox_num']
        start = 0
        for i, im_id in enumerate(outs['im_id']):
            image_path = imid2path[int(im_id)]
            image = Image.open(image_path).convert('RGB')
            image = ImageOps.exif_transpose(image)
            status['original_image'] = np.array(image.copy())
            end = start + bbox_num[i]
            bbox_res = batch_res['bbox'][start:end] \
                if 'bbox' in batch_res else None

            save_label_name = _get_save_image_name(labels_dir, image_path)

            # 通过检测结果生成标签
            getImagesLabels(image, bbox_res, int(im_id), save_label_name, catid2name, draw_threshold)

            # 可视化检测结果
            image = visualize_results(
                image, bbox_res, None, None, None,
                int(im_id), catid2name, draw_threshold)
            status['result_image'] = np.array(image.copy())
            # save image with detection
            save_name = _get_save_image_name(output_dir, image_path)
            # print("save_name: ", save_name)
            logger.info("Detection bbox results save in {}".format(
                save_name))
            image.save(save_name, quality=95)
            start = end


def main():
    FLAGS = parse_args()
    cfg = load_config(FLAGS.config)
    cfg['use_vdl'] = FLAGS.use_vdl
    cfg['vdl_log_dir'] = FLAGS.vdl_log_dir
    merge_config(FLAGS.opt)

    # disable npu in config by default
    if 'use_npu' not in cfg:
        cfg.use_npu = False

    if cfg.use_gpu:
        place = paddle.set_device('gpu')
    elif cfg.use_npu:
        place = paddle.set_device('npu')
    else:
        place = paddle.set_device('cpu')

    if 'norm_type' in cfg and cfg['norm_type'] == 'sync_bn' and not cfg.use_gpu:
        cfg['norm_type'] = 'bn'

    if FLAGS.slim_config:
        cfg = build_slim_model(cfg, FLAGS.slim_config, mode='test')

    check_config(cfg)
    check_gpu(cfg.use_gpu)
    check_npu(cfg.use_npu)
    check_version()

    run(FLAGS, cfg)


if __name__ == '__main__':
    main()