【实例分割】可视化COCO格式Mask分割标注

可视化代码：

import json
import os
import cv2
import numpy as np
from PIL import Image, ImageDraw, ImageColor

def visualise_coco_segmentation(json_file_path, image_dir, output_dir, label_dir, categories, colors):
    category_colors = {category: (color[0], color[1], color[2], 128) for category, color in zip(categories, colors)}

    with open(json_file_path, 'r') as f:
        coco_data = json.load(f)

    images = coco_data['images']
    annotations = coco_data['annotations']
    categories = coco_data['categories']

    for image_info in images:
        image_id = image_info['id']
        image_filename = image_dir + '/' + image_info['file_name']
        image_annotations = [ann for ann in annotations if ann['image_id'] == image_id]
        img = Image.open(image_filename).convert('RGBA')
        overlay = Image.new('RGBA', img.size)
        draw = ImageDraw.Draw(overlay)

        for ann in image_annotations:
            category_id = ann['category_id']
            category = next((cat for cat in categories if cat['id'] == category_id), None)
            if category is None:
                continue

            segmentation = ann['segmentation']
            label = category['name']
            color = category_colors.get(label, (255, 255, 255, 128))

            for segment in segmentation:
                reshaped_polygon = np.array(segment).reshape(-1, 2)
                pil_polygon = [tuple(point) for point in reshaped_polygon]
                draw.polygon(pil_polygon, fill=color)

        img = Image.alpha_composite(img, overlay)
        img = cv2.cvtColor(np.array(img), cv2.COLOR_RGBA2BGRA)

        label_file = f'{label_dir}/{os.path.splitext(image_info["file_name"])[0]}.txt'
        with open(label_file, 'r') as f:
            cnt_info = f.readlines()

        new_cnt_info = [line.replace("\n", "").split(" ") for line in cnt_info]
        height, width, _ = img.shape

        for new_info in new_cnt_info:
            cls_index = int(new_info[0])
            cls = categories[cls_index]['name']
            color = tuple(reversed(colors[cls_index][:3])) + (255,)

            s = [int(float(num) * (width if i % 2 == 0 else height)) for i, num in enumerate(new_info[1:])]
            vertices = np.array(s).reshape(-1, 2)
            x, y, w, h = cv2.boundingRect(vertices)
            # cv2.rectangle(img, (x, y), (x + w, y + h), color, 6)

        output_filename = os.path.splitext(image_info['file_name'])[0] + '.jpg'
        output_path = os.path.join(output_dir, output_filename)
        cv2.imwrite(output_path, img)

        print(f'Image with bounding box processed and saved to {output_path}')

if __name__ == '__main__':
    categories = ['class 1', 'class 2', 'class 3']
    colors = [(255, 0, 217), (39, 26, 213), (140, 224, 56)]
    visualise_coco_segmentation(json_file_path='D:\\coco\\annotations\\instances_test2017.json',
                                image_dir="D:\\YOLO\\images\\test",
                                output_dir='D:\\Ground-truth',
                                label_dir="D:\\YOLO\\labels\\test",
                                categories=categories,
                                colors=colors)