-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathcocodataset_make_mask.py
139 lines (102 loc) · 4.37 KB
/
cocodataset_make_mask.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
# import json
# from PIL import Image, ImageDraw
# import matplotlib.pyplot as plt
# import matplotlib.patches as patches
# import os
# import numpy as np
# # COCOアノテーションJSONファイルのパス
# json_path = "C:/Users/Yasukawa Lab/Desktop/segmantation_tools/hibikino_segmentation.v1i.coco-segmentation/train/_annotations.coco.json"
# with open(json_path, "r") as f:
# coco_data = json.load(f)
# # マスク画像保存用のディレクトリ
# output_dir = "C:/Users/Yasukawa Lab/Desktop/segmantation_tools/hibikino_segmentation.v1i.coco-segmentation/outputs"
# pre_id = 0
# color_map = {
# "ripe_tomato" : {},
# "green_tomato" : {},
# "main_stem" : {},
# "peduncle" : {},
# }
# # 各画像ごとに処理を行う
# for annotation in coco_data["annotations"]:
# # 画像ID
# image_id = annotation["image_id"]
# print(image_id)
# # 画像ファイル名
# image_info = next(item for item in coco_data["images"] if item["id"] == image_id)
# image_file_name = image_info["file_name"]
# # print(image_file_name)
# width = image_info["width"]
# height = image_info["height"]
# # マスク画像の保存パス
# mask_path = os.path.join(output_dir, image_file_name.replace(".jpg", "_mask.png"))
# # マスク画像が存在する場合は読み込み、ない場合は新規作成
# if os.path.exists(mask_path):
# mask = Image.open(mask_path).convert("L")
# else:
# mask = Image.new("L", (width, height), 0)
# draw = ImageDraw.Draw(mask)
# # セグメンテーション情報からポリゴン描画
# for segmentation in annotation["segmentation"]:
# draw.polygon(segmentation, outline=1, fill=1)
# # マスク画像の保存
# mask.save(mask_path)
# # マスク画像を描画
# plt.figure()
# plt.imshow(mask, cmap="gray")
# plt.title("Mask for {}".format(image_file_name))
# plt.show()
# pre_id = image_id
import json
from PIL import Image, ImageDraw
import matplotlib.pyplot as plt
import os
# COCOアノテーションJSONファイルのパス
json_path = "C:/Users/Yasukawa Lab/Desktop/segmantation_tools/hibikino_segmentation.v1i.coco-segmentation/train/_annotations.coco.json"
with open(json_path, "r") as f:
coco_data = json.load(f)
# マスク画像保存用のディレクトリ
output_dir = "C:/Users/Yasukawa Lab/Desktop/segmantation_tools/hibikino_segmentation.v1i.coco-segmentation/outputs"
# 各カテゴリごとに異なる色を割り当てる
# color_map = {
# "ripe_tomato": (255, 0, 0), # Red
# "green_tomato": (0, 255, 0), # Green
# "main_stem": (0, 0, 255), # Blue
# "peduncle": (255, 255, 0), # Yellow
# }
color_map = {
1 : (0, 255, 0), #Green_Tomato #Right_Green
2 : (0, 100, 0), #main_stem #Green
3 : (0, 0, 255), #peduncle #Blue
4 : (255, 0, 0), #ripe_tomato #Red
}
# 各画像ごとに処理を行う
for annotation in coco_data["annotations"]:
# 画像ID
image_id = annotation["image_id"]
# 画像ファイル名
image_info = next(item for item in coco_data["images"] if item["id"] == image_id)
image_file_name = image_info["file_name"]
# print(image_file_name)
width = image_info["width"]
height = image_info["height"]
# マスク画像の保存パス
mask_path = os.path.join(output_dir, image_file_name.replace(".jpg", "_mask.png"))
# マスク画像が存在する場合は読み込み、ない場合は新規作成
if os.path.exists(mask_path):
mask = Image.open(mask_path).convert("RGB")
else:
mask = Image.new("RGB", (width, height), (0, 0, 0))
draw = ImageDraw.Draw(mask)
# セグメンテーション情報からポリゴン描画
for segmentation in annotation["segmentation"]:
category = annotation["category_id"]
color = color_map.get(category, (0, 0, 0)) # マッピングがない場合は黒色を使用
draw.polygon(segmentation, outline=color, fill=color)
# マスク画像の保存
mask.save(mask_path)
# マスク画像を描画
# plt.figure()
# plt.imshow(mask)
# plt.title("Mask for {}".format(image_file_name))
# plt.show()