【基于detectron2训练数据集】

1. 首先下载官方提供的baloon数据集

import wget
from zipfile import ZipFile


def progress_bar(current, total, width=80):
    progress = current / total
    bar = '#' * int(progress * width)
    percentage = round(progress * 100, 2)
    print(f'[{
     
     bar:<{
     
     width}}] {
     
     percentage}%')


save_path = 'balloon_dataset.zip'
if True:
    url = 'https://github.com/matterport/Mask_RCNN/releases/download/v2.1/balloon_dataset.zip'
    try:
        wget.download(url, save_path, bar=progress_bar)
    except Exception as e:
        print(f'An error occurred: {
     
     e}')

extract_path = 'balloon_dataset'
with ZipFile(save_path, "r") as zip:
    zip.printdir()
    zip.extractall(extract_path)

2. 转换成detectron2的格式

# ----------------------------------------------------------------------------
# 转换成 detectron2 的数据格式
# if your dataset is in COCO format, this cell can be replaced by the following three lines:
# from detectron2.data.datasets import register_coco_instances
# register_coco_instances("my_dataset_train", {}, "json_annotation_train.json", "path/to/image/dir")
# register_coco_instances("my_dataset_val", {}, "json_annotation_val.json", "path/to/image/dir")
import cv2
import os
import json
import numpy as np
import random
# import some common detectron2 utilities
from detectron2 import model_zoo
from detectron2.engine import DefaultPredictor
from detectron2.config import get_cfg
from detectron2.utils.visualizer import Visualizer
from detectron2.data import MetadataCatalog, DatasetCatalog
from detectron2.structures import BoxMode


def get_balloon_dicts(img_dir):
    json_file = os.path.join(img_dir, "via_region_data.json")
    with open(json_file) as f:
        imgs_anns = json.load(f)

    dataset_dicts = []
    for idx, v in enumerate(imgs_anns.values()):
        record = {
   
   }

        filename = os.path.join(img_dir, v["filename"])
        height, width = cv2.imread(filename).shape[:2]

        record["file_name"] = filename
        record["image_id"] = idx
        record["height"] = height
        record["width"] = width

        annos = v["regions"]
        objs = []
        for _, anno in annos.items():
            assert not anno["region_attributes"]
            anno = anno["shape_attributes"]
            px = anno["all_points_x"]
            py = anno["all_points_y"]
            poly = [(x + 0.5, y + 0.5