from pycocotools.coco import COCO
import cv2
import numpy as np
import os
from pycocotools import coco
def get_contours(img):
img = cv2.imread(img)
gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
kernel = np.ones((1, 5), np.uint8)
ret, img1 = cv2.threshold(gray, 127, 255, cv2.THRESH_BINARY)
binary = cv2.morphologyEx(img1, cv2.MORPH_CLOSE, kernel, anchor=(2, 0), iterations=5)
contours, hierarchy = cv2.findContours(binary, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
# cv2.drawContours(img1, contours, -1, (0, 0, 134), 3)
return contours, hierarchy
def ReadFromCOCO(path):
# 初始化COCO数据加载器
coco = COCO(path)
# 获取所有图片ID
img_ids = coco.getImgIds()
# 创建空列表用于存储图片路径和标注信息
img_paths_list = []
labels_list = []
# 遍历每个图片ID
for img_id in img_ids:
# 获取图片信息
img_info = coco.loadImgs(img_id)[0]
# 获取图片路径
img_path = img_info['file_name']
# 将图片路径添加到img_paths_list中
img_paths = os.path.join(img_dir, img_path)
img_paths_list.append(img_paths)
# 获取图片对应的标注信息
ann_ids = coco.getAnnIds(imgIds=img_id)
annotations = coco.loadAnns(ann_ids)
# 创建一个空列表用于存储当前图片的标注信息
img_labels = []
for ann in annotations:
# 获取类别ID
category_id = ann['category_id']
# 获取多边形边界信息
segmentation = ann['segmentation']
# 将类别ID和多边形边界信息分别存储在一个小列表中，并添加到img_labels中
img_labels.append([category_id, segmentation])
# 将当前图片的标注信息添加到labels_list中
labels_list.append(img_labels)
return img_paths_list, labels_list
def getMaskFromList(image_path, label):
# 获取指定索引的图片路径
Img_ori, Img_Masks = [], []
print(image_path)
# 读取图片
Img = cv2.imread(image_path)
# # 获取指定索引的标注信息
# print("labels:",label)
# # 创建空列表，用于存储单目标掩膜图像
count = 0
# file1=os.path.join(out_dir,image_path.split("/")[-1].split(".")[0])
# if not os.path.exists(file1): os.makedirs(file1)
for lab in label:
print("label:", label)
category_id, segmentation = lab
print(category_id, segmentation)
if isinstance(segmentation, dict):
img1 = Img.copy()
# 多边形形状，使用RLE解码
rle, m = segmentation["counts"], segmentation["size"]
print(m)
assert sum(rle) == m[0] * m[1]
M = np.zeros(m[0] * m[1])
N = len(rle)
n = 0
val = 1
for pos in range(N):
val = not val
for c in range(rle[pos]):
M[n] = val
if val:
M[n] = 255
n += 1
GEMFIELD = M.reshape(([m[0], m[1]]), order="F")
# file1 = os.path.join(out_dir, image_path.split("/")[-1].split(".")[0])
# if not os.path.exists(file1): os.makedirs(file1)
# file2 = os.path.join(os.path.join(out_dir, image_path.split("/")[-1].split(".")[0]), str(category_id))
# if not os.path.exists(file2): os.makedirs(file2)
# cv2.imwrite(os.path.join(file2, image_path.split("/")[-1].split(".")[0] + "_" + str(count) + ".jpg"),
# GEMFIELD)
Img_Masks.append([category_id, GEMFIELD])
# cv2.imshow('111', GEMFIELD)
# mypath=os.path.join(file2, image_path.split("/")[-1].split(".")[0] + "_" + str(count) + ".jpg")
# contours,hierarchy=get_contours(mypath)
# cv2.drawContours(img1, contours, -1, (255,0,0), 3)
# cv2.imshow("111222",img1)
# cv2.waitKey(0)
# elif isinstance(segmentation, list):
else:
# 列表形式，表示多边形的轮廓点
# print(len(segmentation))
for seg in segmentation:
seg = np.array(seg, dtype=np.int32).reshape(-1, 2)
mask = np.zeros(Img.shape[:2], dtype=np.uint8)
img1 = Img.copy()
file2 = os.path.join(os.path.join(out_dir, image_path.split("/")[-1].split(".")[0]), str(category_id))
if not os.path.exists(file2): os.makedirs(file2)
# cv2.fillPoly(img1, [seg], color=255)
count += 1
# 填充多边形区域
cv2.fillPoly(mask, [seg], color=255)
# cv2.circle(img1, [seg], 5, color=123, thickness=2)
# cv2.imwrite(os.path.join(file2,image_path.split("/")[-1].split(".")[0]+"_"+str(count)+".png"),mask)
# cv2.imwrite(os.path.join(file2, image_path.split("/")[-1].split(".")[0] + "_" + str(count) + ".png"),
# img1)
# cv2.imshow("12345",img1)
# cv2.waitKey(0)
Img_Masks.append([category_id, mask])
return Img, Img_Masks,
# 一张图对应生成一张mask图
# def getMaskFromList1(image_path, labels):
# # 获取指定索引的图片路径
# # 读取图片
# Img = cv2.imread(image_path)
# # 创建空列表，用于存储单目标掩膜图像
# Img_Masks = []
# # mask = np.zeros(Img.shape[:2], dtype=np.uint8)
# # 遍历标注信息，获取单目标掩膜
# for label in labels:
# category_id, segmentation = label
# mask = np.zeros(Img.shape[:2], dtype=np.uint8)
# # 判断segmentation的数据类型，如果是字典，则是多边形形状，否则为列表
# if isinstance(segmentation, dict):
# # 多边形形状，使用RLE解码
# rle, m = segmentation["counts"], segmentation["size"]
# M = np.zeros(m[0] * m[1], dtype=np.uint8)
# N = len(rle)
# n = 0
# val = 255
# for pos in range(N):
# val = not val
# for c in range(rle[pos]):
# M[n] = val
# n += 1
# for idx in range(len(M)):
# if M[idx]==1:
# M[idx]=255
# GEMFIELD = M.reshape(([m[0], m[1]]), order="F")
# cv2.imwrite(os.path.join(out_dir, img_path.split("/")[-1].split(".")[0]+"_1.png"), GEMFIELD)
#
# elif isinstance(segmentation, list):
# # 列表形式，表示多边形的轮廓点
# mask1=[]
# for seg in segmentation:
# seg = np.array(seg, dtype=np.int32).reshape(-1, 2)
# # mask = np.zeros(Img.shape[:2], dtype=np.uint8)
# # 填充多边形区域
# cv2.fillPoly(mask, [seg], color=255)
# Img_Masks.append(mask)
# cv2.imwrite(os.path.join(out_dir, img_path.split("/")[-1]), mask)
#
# return Img, Img_Masks
def get_image_size(img,size):#size=[width,height]
img = np.array(img)
height, width = img.shape
ratio = width / height
mask_image = []
if width > height:
w = size[0]
h = int(w / ratio)
img1 = cv2.resize(img, (w, h))
h1, w1 = img1.shape
if h1 % 2 == 0:
c1 = int((size[1] - h1) / 2)
c2 = int((size[1]- h1) / 2)
mask1 = np.zeros([c1, w], dtype=np.uint8)
mask2 = np.zeros([c2, w], dtype=np.uint8)
# img2=cv2.vconcat([mask1,img1,mask2])
img2 = np.concatenate([mask1, img1, mask2], axis=0)
mask_image.append(img2)
print("*****", img2.shape)
elif h1 % 2 == 1:
c1 = int(((size[1] - h1) / 2) + 1)
c2 = int((size[1] - h1) / 2)
mask1 = np.zeros([c1, w], dtype=np.uint8)
mask2 = np.zeros([c2, w], dtype=np.uint8)
# img2 = cv2.vconcat([mask1, img1, mask2])
img2 = np.concatenate([mask1, img1, mask2], axis=0)
mask_image.append(img2)
elif width < height:
h = size[1]
w = int(h * ratio)
img1 = cv2.resize(img, (w, h))
h1, w1 = img1.shape
if w1 % 2 == 0:
c1 = int((size[0] - w1) / 2)
c2 = int((size[0] - w1) / 2)
mask1 = np.zeros([h, c1], dtype=np.uint8)
mask2 = np.zeros([h, c2], dtype=np.uint8)
# img2 = cv2.hconcat([mask1, img1, mask2])
img2 = np.concatenate([mask1, img1, mask2], axis=1)
mask_image.append(img2)
print(">>>>>>>>>", img2.shape)
# cv2.imwrite("", img2)
elif w1 % 2 == 1:
c1 = int((size[0] - w1) / 2) + 1
c2 = int((size[0] - w1) / 2)
print("c1c2:", c1, c2)
mask1 = np.zeros([h, c1], dtype=np.uint8)
mask2 = np.zeros([h, c2], dtype=np.uint8)
# img2 = cv2.hconcat([mask1, img1, mask2])
img2 = np.concatenate([mask1, img1, mask2], axis=1)
print(">>>>>>>>>", img2.shape)
mask_image.append(img2)
else:
img1 = cv2.resize(img, (size[0], size[1]))
mask_image.append(img1)
return mask_image[0]
if __name__ == '__main__':
data_path = '/media/tianhailong/新加卷/coco_seg/instances_train2017.json'
img_dir = "/media/tianhailong/新加卷/coco_seg/train2017"
out_dir = "/media/tianhailong/新加卷/coco_seg2"
if not os.path.exists(out_dir): os.makedirs(out_dir)
img_paths_list, labels_list = ReadFromCOCO(data_path)
for idx, img_path in enumerate(img_paths_list):
# print(img_path.split("/")[-1])
Img, Img_Masks = getMaskFromList(img_path, labels_list[idx])
file1 = os.path.join(out_dir, img_path.split("/")[-1].split(".")[0])
if not os.path.exists(file1): os.makedirs(file1)
for class_id,mask_img in Img_Masks:
count=0
print(mask_img.shape)
mask_image=get_image_size(mask_img,(640,640))
file2 = os.path.join(os.path.join(out_dir, img_path.split("/")[-1].split(".")[0]), str(class_id))
if not os.path.exists(file2): os.makedirs(file2)
count+=1
cv2.imwrite(os.path.join(file2,img_path.split("/")[-1].split(".")[0]+ "_%s"%(count)+".png"),mask_image)