PaddleX/paddlex_restful/restful/dataset/seg_dataset.py

# copyright (c) 2020 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 os.path as osp
import cv2
from ..utils import list_files
from .utils import is_pic, replace_ext, get_encoding, check_list_txt, read_seg_ann
from .datasetbase import DatasetBase


class SegDataset(DatasetBase):
    def __init__(self, dataset_id, path):
        super().__init__(dataset_id, path)

    def check_dataset(self, source_path):
        if not osp.isdir(osp.join(source_path, 'Annotations')):
            raise ValueError("标注文件应该放在{}目录下".format(
                osp.join(source_path, 'Annotations')))
        if not osp.isdir(osp.join(source_path, 'JPEGImages')):
            raise ValueError("图片文件应该放在{}目录下".format(
                osp.join(source_path, 'JPEGImages')))

        labels_txt = osp.join(source_path, 'labels.txt')
        if osp.exists(labels_txt):
            with open(labels_txt, encoding=get_encoding(labels_txt)) as fid:
                lines = fid.readlines()
                for line in lines:
                    self.labels.append(line.strip())

        self.all_files = list_files(source_path)
        # 对语义分割数据集进行统计分析
        self.file_info = dict()
        self.label_info = dict()

        if osp.exists(osp.join(source_path, 'train_list.txt')):
            return self.check_splited_dataset(source_path)

        for f in self.all_files:
            if not is_pic(f):
                continue
            items = osp.split(f)
            if len(items) == 2 and items[0] == "JPEGImages":
                anno_name = replace_ext(items[1], "png")
                full_anno_path = osp.join(
                    (osp.join(source_path, 'Annotations')), anno_name)
                if osp.exists(full_anno_path):
                    self.file_info[f] = osp.join('Annotations', anno_name)

                # 解析PNG标注文件，获取类别信息
                labels, ann_img_shape = read_seg_ann(full_anno_path)
                img_shape = cv2.imread(osp.join(source_path, f)).shape
                if img_shape[0] != ann_img_shape[0] or img_shape[
                        1] != ann_img_shape[1]:
                    raise ValueError("文件{}与标注图片尺寸不一致".format(items[1]))
                for i in labels:
                    if str(i) not in self.label_info:
                        self.label_info[str(i)] = list()
                    self.label_info[str(i)].append(f)

        # 如果类标签的最大值大于类别数，统计相应的类别为零
        max_label = max([int(i) for i in self.label_info]) + 1
        for i in range(max_label):
            if str(i) not in self.label_info:
                self.label_info[str(i)] = list()

        if len(self.labels) == 0:
            self.labels = [int(i) for i in self.label_info]
            self.labels.sort()
            self.labels = [str(i) for i in self.labels]
        else:
            keys = list(self.label_info.keys())
            try:
                for key in keys:
                    label = self.labels[int(key)]
                    self.label_info[label] = self.label_info.pop(key)
            except:
                raise ValueError("标注信息与实际类别不一致")

        for label in self.labels:
            self.class_train_file_list[label] = list()
            self.class_val_file_list[label] = list()
            self.class_test_file_list[label] = list()

        # 将数据集分析信息dump到本地
        self.dump_statis_info()

    def check_splited_dataset(self, source_path):
        labels_txt = osp.join(source_path, "labels.txt")
        train_list_txt = osp.join(source_path, "train_list.txt")
        val_list_txt = osp.join(source_path, "val_list.txt")
        test_list_txt = osp.join(source_path, "test_list.txt")
        for txt_file in [train_list_txt, val_list_txt]:
            if not osp.exists(txt_file):
                raise Exception("已切分的数据集下应该包含train_list.txt, val_list.txt文件")
        check_list_txt([train_list_txt, val_list_txt, test_list_txt])

        if osp.exists(labels_txt):
            self.labels = open(
                labels_txt, 'r',
                encoding=get_encoding(labels_txt)).read().strip().split('\n')

        for txt_file in [train_list_txt, val_list_txt, test_list_txt]:
            if not osp.exists(txt_file):
                continue
            with open(txt_file, "r") as f:
                for line in f:
                    items = line.strip().split()
                    img_file, png_file = [items[0], items[1]]

                    if not osp.isfile(osp.join(source_path, png_file)):
                        raise ValueError("数据目录{}中不存在标注文件{}".format(
                            osp.split(txt_file)[-1], png_file))
                    if not osp.isfile(osp.join(source_path, img_file)):
                        raise ValueError("数据目录{}中不存在图片文件{}".format(
                            osp.split(txt_file)[-1], img_file))
                    if not png_file.split('.')[-1] == 'png':
                        raise ValueError("标注文件{}不是png文件".format(png_file))
                    img_file_name = osp.split(img_file)[-1]
                    if not is_pic(img_file_name) or img_file_name.startswith(
                            '.'):
                        raise ValueError("文件{}不是图片格式".format(img_file_name))

                    self.file_info[img_file] = png_file

                    if txt_file == train_list_txt:
                        self.train_files.append(img_file)
                    elif txt_file == val_list_txt:
                        self.val_files.append(img_file)
                    elif txt_file == test_list_txt:
                        self.test_files.append(img_file)

                    # 解析PNG标注文件
                    labels, ann_img_shape = read_seg_ann(
                        osp.join(source_path, png_file))
                    img_shape = cv2.imread(osp.join(source_path,
                                                    img_file)).shape
                    if img_shape[0] != ann_img_shape[0] or img_shape[
                            1] != ann_img_shape[1]:
                        raise ValueError("文件{}与标注图片尺寸不一致".format(
                            img_file_name))
                    for i in labels:
                        if str(i) not in self.label_info:
                            self.label_info[str(i)] = list()
                        self.label_info[str(i)].append(img_file)

        # 如果类标签的最大值大于类别数，统计相应的类别为零
        max_label = max([int(i) for i in self.label_info]) + 1
        for i in range(max_label):
            if str(i) not in self.label_info:
                self.label_info[str(i)] = list()

        if len(self.labels) == 0:
            self.labels = [int(i) for i in self.label_info]
            self.labels.sort()
            self.labels = [str(i) for i in self.labels]
        else:
            keys = list(self.label_info.keys())
            try:
                for key in keys:
                    label = self.labels[int(key)]
                    self.label_info[label] = self.label_info.pop(key)
            except:
                raise ValueError("标注信息与实际类别不一致")

        self.train_set = set(self.train_files)
        self.val_set = set(self.val_files)
        self.test_set = set(self.test_files)

        for label, file_list in self.label_info.items():
            self.class_train_file_list[label] = list()
            self.class_val_file_list[label] = list()
            self.class_test_file_list[label] = list()
            for f in file_list:
                if f in self.test_set:
                    self.class_test_file_list[label].append(f)
                if f in self.val_set:
                    self.class_val_file_list[label].append(f)
                if f in self.train_set:
                    self.class_train_file_list[label].append(f)

        # 将数据集分析信息dump到本地
        self.dump_statis_info()

    def split(self, val_split, test_split):
        super().split(val_split, test_split)
        with open(
                osp.join(self.path, 'train_list.txt'), mode='w',
                encoding='utf-8') as f:
            for x in self.train_files:
                label = self.file_info[x]
                f.write('{} {}\n'.format(x, label))
        with open(
                osp.join(self.path, 'val_list.txt'), mode='w',
                encoding='utf-8') as f:
            for x in self.val_files:
                label = self.file_info[x]
                f.write('{} {}\n'.format(x, label))
        with open(
                osp.join(self.path, 'test_list.txt'), mode='w',
                encoding='utf-8') as f:
            for x in self.test_files:
                label = self.file_info[x]
                f.write('{} {}\n'.format(x, label))
        if not osp.exists(osp.join(self.path, 'labels.txt')):
            with open(
                    osp.join(self.path, 'labels.txt'), mode='w',
                    encoding='utf-8') as f:
                max_label = max([int(i) for i in self.labels]) + 1
                for i in range(max_label):
                    f.write('{}\n'.format(str(i)))
        self.dump_statis_info()