Блог

Мотивация

• У меня есть базовая модель detectron2 Mask R-CNN, которая достаточно хороша, чтобы точно предсказать границы некоторых объектов.
• Я хотел бы преобразовать эти прогнозируемые границы в кокосовые полигоны для аннотирования следующего набора данных (контролируемая маркировка).
• Для этого мне нужно выполнить вывод для набора данных изображений, у которого нет аннотаций.
• Методы register_coco_instances detectron2 и load_coco_json требуют изображений с аннотациями для правильной маркировки предсказанных объектов.

Вопросы

• Могу ли я зарегистрировать тестовый набор данных без файла аннотаций?
• Если нет, какой самый простой способ сгенерировать файлы COCO или Labelme JSON с базовой информацией об изображении без аннотаций?

Код

 dataset_name = "test_data"
image_dir = "data/test"
coco_file = "data/test_annotations.json"

# Register dataset
# A COCO file is needed with image info, which I don't have
register_coco_instances(dataset_name , {}, coco_file, image_dir)
test_dict = load_coco_json(coco_file, image_dir, dataset_name=dataset_name )
metadata = MetadataCatalog.get(dataset_name)

# config details omitted for brevity
cfg = get_cfg()
predictor = DefaultPredictor(cfg)

# Make predictions for all images
for sample in test_dict:
    image_filename = sample["file_name"]
    img = cv2.imread(image_filename)
    outputs = predictor(img)
    # Display or save image with predictions to file
 

Вот метод для генерации деталей изображения из каталога изображений и записи его в существующий файл COCO JSON:

 from PIL import Image

def add_images_to_coco(image_dir, coco_filename):
    image_filenames = list(Path(image_dir).glob('*.jpg'))
    images = []
    for i, image_filename in enumerate(image_filenames):
        im = Image.open(image_filename)
        width, height = im.size
        image_details = {
            "id": id   1,
            "height": height,
            "width": width,
            "file_name": str(image_filename.resolve()),
        }
        images.append(image_details)

    # This will overwrite the image tags in the COCO JSON file
    with open(coco_filename) as f:
        data = json.load(f)

    coco['images'] = images

    with open(coco_filename, 'w') as coco_file:
        json.dump(data, coco_file, indent = 4)
 

Вам нужно будет создать базовый файл COCO JSON с вашими категориями, если у вас его еще нет. Это должно выглядеть примерно так:

 {
    "images": [ ],
    "annotations": [ ],
    "categories": [
        {
            "id": 1,
            "name": "Car",
            "supercategory": "Car"
        },
        {
            "id": 2,
            "name": "Person",
            "supercategory": "Person"
        }
    ]
}