<reponame>liuyanyi/mmdetection<filename>tools/dota/dota_image.py
import os
from argparse import ArgumentParser

from mmdet.apis import (inference_detector,
                        init_detector)


def parse_args():
    parser = ArgumentParser()
    parser.add_argument('img_dir', help='Image file folder')
    parser.add_argument('out_dir', help='Image file output folder')
    parser.add_argument('config', help='Config file')
    parser.add_argument('checkpoint', help='Checkpoint file')
    parser.add_argument(
        '--device', default='cuda:0', help='Device used for inference')
    parser.add_argument(
        '--score-thr', type=float, default=0.3, help='bbox score threshold')
    args = parser.parse_args()
    return args


def main(args):
    # build the model from a config file and a checkpoint file
    model = init_detector(args.config, args.checkpoint, device=args.device)

    files = os.listdir(args.img_dir)  # 得到文件夹下的所有文件名称
    for file in files:  # 遍历文件夹
        if not os.path.isdir(file):  # 判断是否是文件夹，不是文件夹才打开
            file_path = os.path.join(args.img_dir, file)
            # test a single image
            result = inference_detector(model, file_path)

            # show the results
            model.show_result(
                file_path,
                result,
                score_thr=args.score_thr,
                show=False,
                wait_time=0,
                bbox_color=(72, 101, 241),
                text_color=(72, 101, 241),
                out_file=os.path.join(args.out_dir, file)
            )


if __name__ == '__main__':
    args = parse_args()
    main(args)

import os
from argparse import ArgumentParser

from mmdet.apis import (inference_detector,
                        init_detector)


def parse_args():
    parser = ArgumentParser()
    parser.add_argument('img_dir', help='Image file folder')
    parser.add_argument('out_dir', help='Image file output folder')
    parser.add_argument('config', help='Config file')
    parser.add_argument('checkpoint', help='Checkpoint file')
    parser.add_argument(
        '--device', default='cuda:0', help='Device used for inference')
    parser.add_argument(
        '--score-thr', type=float, default=0.3, help='bbox score threshold')
    args = parser.parse_args()
    return args


def main(args):
    # build the model from a config file and a checkpoint file
    model = init_detector(args.config, args.checkpoint, device=args.device)

    files = os.listdir(args.img_dir)  # 得到文件夹下的所有文件名称
    for file in files:  # 遍历文件夹
        if not os.path.isdir(file):  # 判断是否是文件夹，不是文件夹才打开
            file_path = os.path.join(args.img_dir, file)
            # test a single image
            result = inference_detector(model, file_path)

            # show the results
            model.show_result(
                file_path,
                result,
                score_thr=args.score_thr,
                show=False,
                wait_time=0,
                bbox_color=(72, 101, 241),
                text_color=(72, 101, 241),
                out_file=os.path.join(args.out_dir, file)
            )


if __name__ == '__main__':
    args = parse_args()
    main(args)
# build the model from a config file and a checkpoint file # 得到文件夹下的所有文件名称 # 遍历文件夹 # 判断是否是文件夹，不是文件夹才打开 # test a single image # show the results