概要
在任务过程中缺少必要的数据,且该类型数据不易捕获。只能通过后期合成的方式来得到数据。利用COCO的分割数据集(80个类),得到想要的物体素材,然后通过一系列策略拼接到图像上。
整体架构流程
1.明确COCO的标注构成和标注的格式。
coco数据集是以一个个json文件保存标签。instances_train2017.json就包含了所有train的图像的标签和图像信息。
具体包含的内容如下:
我们主要关注三块,“images”,“annotations”,“categories”。
-
“images”
这里主要包含了图像的名称,长宽,下载的网址链接,coco编号id,这个id是用于映射到相对应的标签上的。同一个image和annotion拥有相同的“id”。
-
“annotations”
包含分割和检测的坐标,还有对应的类别id,图像id等。
3.“categories”
这里包含了类别id和对应中文名称的映射。可以看到,id中间缺少26,所以,类别序号是不连续的,不是0-79. “supercategory”是大类,“name”是大类里面的小类别。我们常说的80类是指小类。
-
正式开始的步骤: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)
- 效果展示
这里只扣取了人的图像,保存格式是png,背景是透明的,可以拼接到其他图像上。