PaddleX/paddlex/modules/anomaly_detection/dataset_checker/dataset_src/convert_dataset.py

# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
#
# 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 json
import os
import os.path as osp
import shutil

import numpy as np
from PIL import Image, ImageDraw

from .....utils import logging
from .....utils.deps import function_requires_deps, is_dep_available
from .....utils.file_interface import custom_open
from .....utils.logging import info

if is_dep_available("opencv-contrib-python"):
    import cv2


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)


@function_requires_deps("opencv-contrib-python")
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}")