利用coco分割数据得到感兴趣的素材

概要

在任务过程中缺少必要的数据，且该类型数据不易捕获。只能通过后期合成的方式来得到数据。利用COCO的分割数据集（80个类），得到想要的物体素材，然后通过一系列策略拼接到图像上。

整体架构流程

1.明确COCO的标注构成和标注的格式。

coco数据集是以一个个json文件保存标签。instances_train2017.json就包含了所有train的图像的标签和图像信息。
具体包含的内容如下：

我们主要关注三块，“images”，“annotations”，“categories”。

1. “images”
   这里主要包含了图像的名称，长宽，下载的网址链接，coco编号id，这个id是用于映射到相对应的标签上的。同一个image和annotion拥有相同的“id”。
   

2. “annotations”
   包含分割和检测的坐标，还有对应的类别id，图像id等。
   
   3.“categories”
   这里包含了类别id和对应中文名称的映射。可以看到，id中间缺少26，所以，类别序号是不连续的，不是0-79. “supercategory”是大类，“name”是大类里面的小类别。我们常说的80类是指小类。
   

3. 正式开始的步骤：1. 转换coco的json标注到txt文件；2.把txt保存的segmentation信息反应到原图上进行裁剪。
   步骤1：

import os
import json
from tqdm import tqdm
import argparse
 
parser = argparse.ArgumentParser()
#这里根据自己的json文件位置，换成自己的就行
parser.add_argument('--json_path', default='C:/Users/Administrator/Desktop/coco/annotations/instances_train2017.json',type=str, help="input: coco format(json)")
#这里设置.txt文件保存位置，一张图片对应生成一个txt
parser.add_argument('--save_path', default='E:/coco/txt/', type=str, help="specify where to save the output dir of labels")
arg = parser.parse_args()
 
if __name__ == '__main__':
    json_file =   arg.json_path # COCO Object Instance 类型的标注
    ana_txt_save_path = arg.save_path  # 保存的路径
 
    data = json.load(open(json_file, 'r'))
    if not os.path.exists(ana_txt_save_path):
        os.makedirs(ana_txt_save_path)
 
    id_map = {} # coco数据集的id不连续！重新映射一下再输出！
    with open(os.path.join(ana_txt_save_path, 'classes.txt'), 'w') as f:
        # 写入classes.txt
        for i, category in enumerate(data['categories']):
            f.write(f"{category['name']}\n")
            id_map[category['id']] = i
    # print(id_map)

    for img in tqdm(data['images']):
        filename = img["file_name"]
        img_width = img["width"]
        img_height = img["height"]
        img_id = img["id"]
        head, tail = os.path.splitext(filename)
        ana_txt_name = head + ".txt"  # 对应的txt名字，与jpg一致
        f_txt = open(os.path.join(ana_txt_save_path, ana_txt_name), 'w')
        
        for ann in data['annotations']:
            if ann['image_id'] == img_id:
                f_txt.write(str(id_map[ann["category_id"]])+" ")
                if len(ann['segmentation'])<=6: #判断有效点小于6的话，就是无效的标签
                    pass
                else:
                    for i in range(len(ann['segmentation'][0])):
                        if i ==len(ann['segmentation'][0])-1:
                            f_txt.write(str(ann['segmentation'][0][i])+"\n")
                        else:
                            f_txt.write(str(ann['segmentation'][0][i])+" ")

        f_txt.close()
print('finish it')

步骤2：

import numpy as np
import cv2
import os

#分割的txt存放地址
dir="E:/coco/txt/"
#要切割的图像地址
img_dir="E:/coco/train2017/"
#保存切割后的图像地址
save_img="E:/coco/save_img/"


for file in os.listdir(dir):
    orig_img = cv2.imread(img_dir+file.replace(".txt",".jpg"))
    canvas_height, canvas_width=orig_img.shape[0], orig_img.shape[1]
    with open(dir+file,"r") as f:
        for line in f.readlines():  
            a=[]
            points=line.split(" ")[1:]
            if len(points)<15:
                continue
            cls=int(line.split(" ")[0])
            if cls==0:
                for i in range(len(points)):
                    a.append(int(points[i].split(".")[0]))

                if len(a)%2==1:
                    continue
                a = np.array(a).reshape(-1, 2) 

                # 在布尔数组上填充多边形区域
                mask = np.zeros((canvas_height, canvas_width))
                cv2.fillPoly(mask, [a], 1)  # True表示填充为True

                # 将矩阵转换为布尔类型的矩阵
                mask = (mask == 1)
                
                # # 如果你想要背景是黑色的
                masked_image = np.zeros_like(orig_img)
                masked_image[mask] = orig_img[mask]

                # 如果你想要背景透明（假设原始图像是RGB格式）
                # 创建一个RGBA图像，其中背景是透明的
                alpha_channel = np.ones(mask.shape, dtype=orig_img.dtype) * 255  # 创建alpha通道
                masked_image_rgba = np.dstack((masked_image, alpha_channel))  # 将alpha通道添加到RGB通道
                masked_image_rgba[~mask] = (0, 0, 0, 0)  # 将非掩码区域设置为透明
                
                # 如果背景是透明的
                cv2.imwrite(save_img+file.replace(".txt",".png"), masked_image_rgba)

3. 效果展示
   这里只扣取了人的图像，保存格式是png，背景是透明的，可以拼接到其他图像上。