|
|
@@ -0,0 +1,210 @@
|
|
|
+# copyright (c) 2024 PaddlePaddle Authors. All Rights Reserve.
|
|
|
+#
|
|
|
+# Licensed under the Apache License, Version 2.0 (the "License");
|
|
|
+# you may not use this file except in compliance with the License.
|
|
|
+# You may obtain a copy of the License at
|
|
|
+#
|
|
|
+# http://www.apache.org/licenses/LICENSE-2.0
|
|
|
+#
|
|
|
+# Unless required by applicable law or agreed to in writing, software
|
|
|
+# distributed under the License is distributed on an "AS IS" BASIS,
|
|
|
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
|
+# See the License for the specific language governing permissions and
|
|
|
+# limitations under the License.
|
|
|
+
|
|
|
+
|
|
|
+import glob
|
|
|
+import json
|
|
|
+import os
|
|
|
+import os.path as osp
|
|
|
+import shutil
|
|
|
+
|
|
|
+import cv2
|
|
|
+import numpy as np
|
|
|
+from PIL import Image, ImageDraw
|
|
|
+
|
|
|
+from .....utils.file_interface import custom_open
|
|
|
+from .....utils import logging
|
|
|
+from .....utils.logging import info
|
|
|
+
|
|
|
+def convert_dataset(dataset_type, input_dir):
|
|
|
+ """convert to paddlex official format"""
|
|
|
+ if dataset_type == "LabelMe":
|
|
|
+ return convert_labelme_dataset(input_dir)
|
|
|
+ elif dataset_type == "MVTec_AD":
|
|
|
+ return convert_mvtec_dataset(input_dir)
|
|
|
+ else:
|
|
|
+ raise NotImplementedError(dataset_type)
|
|
|
+
|
|
|
+
|
|
|
+def convert_labelme_dataset(input_dir):
|
|
|
+ """convert labelme format to paddlex official format"""
|
|
|
+ bg_name = "_background_"
|
|
|
+ ignore_name = "__ignore__"
|
|
|
+
|
|
|
+ # prepare dir
|
|
|
+ output_img_dir = osp.join(input_dir, "images")
|
|
|
+ output_annot_dir = osp.join(input_dir, "annotations")
|
|
|
+ if not osp.exists(output_img_dir):
|
|
|
+ os.makedirs(output_img_dir)
|
|
|
+ if not osp.exists(output_annot_dir):
|
|
|
+ os.makedirs(output_annot_dir)
|
|
|
+
|
|
|
+ # collect class_names and set class_name_to_id
|
|
|
+ class_names = []
|
|
|
+ class_name_to_id = {}
|
|
|
+ split_tags = ["train", "val"]
|
|
|
+ for tag in split_tags:
|
|
|
+ mapping_file = osp.join(input_dir, f"{tag}_anno_list.txt")
|
|
|
+ with open(mapping_file, "r") as f:
|
|
|
+ label_files = [
|
|
|
+ osp.join(input_dir, line.strip("\n")) for line in f.readlines()
|
|
|
+ ]
|
|
|
+ for label_file in label_files:
|
|
|
+ with custom_open(label_file, "r") as fp:
|
|
|
+ data = json.load(fp)
|
|
|
+ for shape in data["shapes"]:
|
|
|
+ cls_name = shape["label"]
|
|
|
+ if cls_name not in class_names:
|
|
|
+ class_names.append(cls_name)
|
|
|
+
|
|
|
+ if ignore_name in class_names:
|
|
|
+ class_name_to_id[ignore_name] = 255
|
|
|
+ class_names.remove(ignore_name)
|
|
|
+ if bg_name in class_names:
|
|
|
+ class_names.remove(bg_name)
|
|
|
+ class_name_to_id[bg_name] = 0
|
|
|
+ for i, name in enumerate(class_names):
|
|
|
+ class_name_to_id[name] = i + 1
|
|
|
+
|
|
|
+ if len(class_names) > 256:
|
|
|
+ raise ValueError(
|
|
|
+ f"There are {len(class_names)} categories in the annotation file, "
|
|
|
+ f"exceeding 256, Not compliant with paddlex official format!"
|
|
|
+ )
|
|
|
+
|
|
|
+ # create annotated images and copy origin images
|
|
|
+ color_map = get_color_map_list(256)
|
|
|
+ img_file_list = []
|
|
|
+ label_file_list = []
|
|
|
+ for i, label_file in enumerate(label_files):
|
|
|
+ filename = osp.splitext(osp.basename(label_file))[0]
|
|
|
+ annotated_img_path = osp.join(output_annot_dir, filename + ".png")
|
|
|
+ with custom_open(label_file, "r") as f:
|
|
|
+ data = json.load(f)
|
|
|
+ img_path = osp.join(osp.dirname(label_file), data["imagePath"])
|
|
|
+ if not os.path.exists(img_path):
|
|
|
+ logging.info("%s is not existed, skip this image" % img_path)
|
|
|
+ continue
|
|
|
+ img_name = img_path.split("/")[-1]
|
|
|
+ img_file_list.append(f"images/{img_name}")
|
|
|
+ label_img_name = annotated_img_path.split("/")[-1]
|
|
|
+ label_file_list.append(f"annotations/{label_img_name}")
|
|
|
+
|
|
|
+ img = np.asarray(cv2.imread(img_path))
|
|
|
+ lbl = shape2label(
|
|
|
+ img_size=img.shape,
|
|
|
+ shapes=data["shapes"],
|
|
|
+ class_name_mapping=class_name_to_id,
|
|
|
+ )
|
|
|
+ lbl_pil = Image.fromarray(lbl.astype(np.uint8), mode="P")
|
|
|
+ lbl_pil.putpalette(color_map)
|
|
|
+ lbl_pil.save(annotated_img_path)
|
|
|
+
|
|
|
+ shutil.copy(img_path, output_img_dir)
|
|
|
+ with custom_open(osp.join(input_dir, f"{tag}.txt"), "w") as fp:
|
|
|
+ for img_path, lbl_path in zip(img_file_list, label_file_list):
|
|
|
+ fp.write(f"{img_path} {lbl_path}\n")
|
|
|
+
|
|
|
+ with custom_open(osp.join(input_dir, "class_name.txt"), "w") as fp:
|
|
|
+ for name in class_names:
|
|
|
+ fp.write(f"{name}{os.linesep}")
|
|
|
+ with custom_open(osp.join(input_dir, "class_name_to_id.txt"), "w") as fp:
|
|
|
+ for key, val in class_name_to_id.items():
|
|
|
+ fp.write(f"{val}: {key}{os.linesep}")
|
|
|
+
|
|
|
+ return input_dir
|
|
|
+
|
|
|
+
|
|
|
+def get_color_map_list(num_classes):
|
|
|
+ """get color map list"""
|
|
|
+ num_classes += 1
|
|
|
+ color_map = num_classes * [0, 0, 0]
|
|
|
+ for i in range(0, num_classes):
|
|
|
+ j = 0
|
|
|
+ lab = i
|
|
|
+ while lab:
|
|
|
+ color_map[i * 3] |= ((lab >> 0) & 1) << (7 - j)
|
|
|
+ color_map[i * 3 + 1] |= ((lab >> 1) & 1) << (7 - j)
|
|
|
+ color_map[i * 3 + 2] |= ((lab >> 2) & 1) << (7 - j)
|
|
|
+ j += 1
|
|
|
+ lab >>= 3
|
|
|
+ color_map = color_map[3:]
|
|
|
+ return color_map
|
|
|
+
|
|
|
+
|
|
|
+def shape2label(img_size, shapes, class_name_mapping):
|
|
|
+ """根据输入的形状列表,将图像的标签矩阵填充为对应形状的类别编号"""
|
|
|
+ label = np.zeros(img_size[:2], dtype=np.int32)
|
|
|
+ for shape in shapes:
|
|
|
+ points = shape["points"]
|
|
|
+ class_name = shape["label"]
|
|
|
+ label_mask = polygon2mask(img_size[:2], points)
|
|
|
+ label[label_mask] = class_name_mapping[class_name]
|
|
|
+ return label
|
|
|
+
|
|
|
+
|
|
|
+def polygon2mask(img_size, points):
|
|
|
+ """将给定形状的点转换成对应的掩膜"""
|
|
|
+ label_mask = Image.fromarray(np.zeros(img_size[:2], dtype=np.uint8))
|
|
|
+ image_draw = ImageDraw.Draw(label_mask)
|
|
|
+ points_list = [tuple(point) for point in points]
|
|
|
+ assert len(points_list) > 2, ValueError("Polygon must have points more than 2")
|
|
|
+ image_draw.polygon(xy=points_list, outline=1, fill=1)
|
|
|
+ return np.array(label_mask, dtype=bool)
|
|
|
+
|
|
|
+
|
|
|
+def save_item_to_txt(items, file_path):
|
|
|
+ try:
|
|
|
+ with open(file_path, 'a') as file:
|
|
|
+ file.write(items)
|
|
|
+ file.close()
|
|
|
+ except Exception as e:
|
|
|
+ print(f"Saving_error: {e}")
|
|
|
+
|
|
|
+
|
|
|
+def save_training_txt(cls_root, mode, cat):
|
|
|
+ imgs = os.listdir(os.path.join(cls_root, mode, cat))
|
|
|
+ imgs.sort()
|
|
|
+ for img in imgs:
|
|
|
+ if mode == 'train':
|
|
|
+ item = os.path.join(cls_root, mode, cat, img)
|
|
|
+ items = item + ' ' + item + '\n'
|
|
|
+ save_item_to_txt(items, os.path.join(cls_root, 'train.txt'))
|
|
|
+ elif mode == 'test' and cat != 'good':
|
|
|
+ item1 = os.path.join(cls_root, mode, cat, img)
|
|
|
+ item2 = os.path.join(cls_root, 'ground_truth', cat, img.split('.')[0]+'_mask.png')
|
|
|
+ items = item1 + ' ' + item2 + '\n'
|
|
|
+ save_item_to_txt(items, os.path.join(cls_root, 'val.txt'))
|
|
|
+
|
|
|
+
|
|
|
+def check_old_txt(cls_pth, mode):
|
|
|
+ set_name = 'train.txt' if mode == 'train' else 'val.txt'
|
|
|
+ pth = os.path.join(cls_pth, set_name)
|
|
|
+ if os.path.exists(pth):
|
|
|
+ os.remove(pth)
|
|
|
+
|
|
|
+
|
|
|
+def convert_mvtec_dataset(input_dir):
|
|
|
+ classes = ['bottle', 'cable', 'capsule', 'hazelnut', 'metal_nut', 'pill', 'screw',
|
|
|
+ 'toothbrush', 'transistor', 'zipper', 'carpet', 'grid', 'leather', 'tile', 'wood']
|
|
|
+ clas = os.path.split(input_dir)[-1]
|
|
|
+ assert clas in classes, info(f"Make sure your class: '{clas}' in your dataset root in\n {classes}")
|
|
|
+ modes = ['train', 'test']
|
|
|
+ cls_root = input_dir
|
|
|
+ for mode in modes:
|
|
|
+ check_old_txt(cls_root, mode)
|
|
|
+ cats = os.listdir(os.path.join(cls_root, mode))
|
|
|
+ for cat in cats:
|
|
|
+ save_training_txt(cls_root, mode, cat)
|
|
|
+ info(f"Add train.txt/val.txt successfully for {input_dir}")
|