x2m.py

import glob
import os
import xml.etree.ElementTree as ET

import numpy as np
import yaml
from mltools.src.log.logger import logger
from mltools.src.utils.xml2yolo.xml2yolo import read_labels
from skimage import io
from tqdm import tqdm


def yaml2dict(filepath):
    if os.path.exists(filepath):
        f = open(filepath)
        y = yaml.load(f, Loader=yaml.FullLoader)
        try:
            f.close()
        except:
            pass
        return y
    else:
        raise FileExistsError("file not found")


def labels2yaml(labels: list, savePath="", savefile=True):
    # pass
    if "_background_" not in labels:
        labels.append("_background_")

    tmp = dict()
    tmp["_background_"] = 0
    classId = 1
    for i in range(0, len(labels)):
        x = labels[i].replace("\n", "").strip()
        if x != "_background_":
            tmp[x] = classId
            classId += 1
    data = dict()
    data["label_names"] = tmp
    del tmp

    if savefile:
        try:
            with open(savePath + os.sep + "info.yaml", "w", encoding="utf-8") as f:
                yaml.dump(data, f, allow_unicode=True)
            logger.info(
                "successfully convert labels to yaml, see {}".format(
                    savePath + os.sep + "info.yaml"
                )
            )
        except Exception as e:
            logger.error(e)

    return data


def generate_mask(xmlPath, parent_path="", label_masks=None, flag=True):
    # classSet = set(labels)
    if parent_path == "":
        parent_path = os.path.dirname(xmlPath)
    fileName = xmlPath.split(os.sep)[-1].replace(".xml", "")
    # out_file = parent_path+os.sep+fileName+'.jpg'

    in_file = open(xmlPath, encoding="utf-8")
    tree = ET.parse(in_file)
    root = tree.getroot()
    size = root.find("size")
    w = int(size.find("width").text)
    h = int(size.find("height").text)

    mask_img = np.zeros((h, w)).astype(np.uint8)

    generateClass = dict()

    for obj in root.iter("object"):
        difficult = obj.find("difficult").text
        clas = obj.find("name").text
        xmlbox = obj.find("bndbox")
        b = (
            int(xmlbox.find("xmin").text),
            int(xmlbox.find("xmax").text),
            int(xmlbox.find("ymin").text),
            int(xmlbox.find("ymax").text),
        )
        # print(b)
        if label_masks is not None:
            try:
                generateClass = label_masks["label_names"]
                classId = generateClass.get(clas)

                if classId is not None:
                    mask_img[b[2] : b[3], b[0] : b[1]] = int(classId)

                    # mask_img[b[2]+15:b[3]-15,b[0]+15:b[1]-15] = int(classId)
                else:
                    vals = generateClass.values()
                    classId = max(vals) + 1
                    mask_img[b[2] : b[3], b[0] : b[1]] = int(classId)
                    generateClass[clas] = classId

                    label_masks["label_names"] = generateClass

            except Exception as e:
                logger.error(e)

        else:
            label_masks = dict()
            generateClass[clas] = 1
            generateClass["_background_"] = 0
            label_masks["label_names"] = generateClass
            mask_img[b[2] : b[3], b[0] : b[1]] = 1
            label_masks["label_names"] = generateClass

    if flag:
        if not os.path.exists(parent_path + os.sep + "mask_"):
            os.makedirs(parent_path + os.sep + "mask_")
        io.imsave(parent_path + os.sep + "mask_" + os.sep + fileName + ".jpg", mask_img)
        # io.imsave(parent_path+os.sep+'mask_'+os.sep+fileName+'.jpg',np.array(mask_img/np.max(mask_img)*255,dtype=np.uint8))
        print(
            "process finished. see {}".format(
                parent_path + os.sep + "mask_" + os.sep + fileName + ".jpg"
            )
        )
        return label_masks, parent_path + os.sep + "mask_" + os.sep + fileName + ".jpg"
    else:
        return label_masks


def x2m_convert(xmlpath, labelPath="", yamlPath=""):
    """this function is used to convert xml to masks"""
    flagY = True
    flagL = True
    if not os.path.exists(yamlPath):
        logger.info("yaml file not exists!")
        flagY = False
        if not os.path.exists(labelPath):
            logger.info("label file not exists!")
            flagL = False
        else:
            labels = read_labels(labelPath)
            label_masks = labels2yaml(labels, savefile=False)
    else:
        label_masks = yaml2dict(yamlPath)

    parent_path = os.path.dirname(xmlpath)
    if not os.path.exists(xmlpath):
        raise FileNotFoundError("file not found")
    else:
        if os.path.isfile(xmlpath):
            logger.info("single file found")

            if flagY or flagL:
                label_masks, maskPath = generate_mask(xmlpath, parent_path, label_masks)
            else:
                label_masks, maskPath = generate_mask(xmlpath, parent_path)
            print("Done!")

            return label_masks, maskPath

        else:
            xmls = glob.glob(xmlpath + os.sep + "*.xml")
            if not os.path.exists(parent_path + os.sep + "masks_"):
                os.mkdir(parent_path + os.sep + "masks_")
            logger.info("exists {} xml files".format(len(xmls)))

            for xml in tqdm(xmls):
                if flagY or flagL:
                    label_masks = generate_mask(xml, parent_path, label_masks)
                else:
                    label_masks = generate_mask(xml, parent_path)

            print("Done!")
            print("see here {}".format(parent_path + os.sep + "masks_"))

            return label_masks, parent_path + os.sep + "masks_"