PaddleX/paddlex_restful/restful/dataset/ins_seg_dataset.py

# copyright (c) 2021 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 random
from ..utils import list_files
from .utils import is_pic, replace_ext, MyEncoder, read_coco_ann, get_npy_from_coco_json
from .datasetbase import DatasetBase
import json
from pycocotools.coco import COCO


class InsSegDataset(DatasetBase):
    def __init__(self, dataset_id, path):
        super().__init__(dataset_id, path)
        self.annotation_dict = None

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

        self.all_files = list_files(source_path)
        # 对检测数据集进行统计分析
        self.file_info = dict()
        self.label_info = dict()
        # 若数据集已切分
        if osp.exists(osp.join(source_path, 'train.json')):
            return self.check_splited_dataset(source_path)

        if not osp.exists(osp.join(source_path, 'annotations.json')):
            raise ValueError("标注文件annotations.json应该放在{}目录下".format(
                source_path))

        filename_set = set()
        anno_set = set()
        for f in self.all_files:
            items = osp.split(f)
            if len(items) == 2 and items[0] == "JPEGImages":
                if not is_pic(f) or f.startswith('.'):
                    continue
                filename_set.add(items[1])
        # 解析包含标注信息的json文件
        try:
            coco = COCO(osp.join(source_path, 'annotations.json'))
            img_ids = coco.getImgIds()
            cat_ids = coco.getCatIds()
            anno_ids = coco.getAnnIds()

            catid2clsid = dict({catid: i for i, catid in enumerate(cat_ids)})
            cid2cname = dict({
                clsid: coco.loadCats(catid)[0]['name']
                for catid, clsid in catid2clsid.items()
            })
            for img_id in img_ids:
                img_anno = coco.loadImgs(img_id)[0]
                img_name = osp.split(img_anno['file_name'])[-1]
                anno_set.add(img_name)
                anno_dict = read_coco_ann(img_id, coco, cid2cname, catid2clsid)
                img_path = osp.join("JPEGImages", img_name)
                anno_path = osp.join("Annotations", img_name)
                anno = replace_ext(anno_path, "npy")
                self.file_info[img_path] = anno

                img_class = list(set(anno_dict["gt_class"]))
                for category_name in img_class:
                    if not category_name in self.label_info:
                        self.label_info[category_name] = [img_path]
                    else:
                        self.label_info[category_name].append(img_path)

            for label in sorted(self.label_info.keys()):
                self.labels.append(label)
        except:
            raise Exception("标注文件存在错误")

        if len(anno_set) > len(filename_set):
            sub_list = list(anno_set - filename_set)
            raise Exception("标注文件中{}等{}个信息无对应图片".format(sub_list[0],
                                                        len(sub_list)))

        # 生成每个图片对应的标注信息npy文件
        npy_path = osp.join(self.path, "Annotations")
        get_npy_from_coco_json(coco, npy_path, self.file_info)

        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):
        train_files_json = osp.join(source_path, "train.json")
        val_files_json = osp.join(source_path, "val.json")
        test_files_json = osp.join(source_path, "test.json")

        for json_file in [train_files_json, val_files_json]:
            if not osp.exists(json_file):
                raise Exception("已切分的数据集下应该包含train.json, val.json文件")

        filename_set = set()
        anno_set = set()
        # 获取全部图片名称
        for f in self.all_files:
            items = osp.split(f)
            if len(items) == 2 and items[0] == "JPEGImages":
                if not is_pic(f) or f.startswith('.'):
                    continue
                filename_set.add(items[1])

        img_id_index = 0
        anno_id_index = 0
        new_img_list = list()
        new_cat_list = list()
        new_anno_list = list()
        for json_file in [train_files_json, val_files_json, test_files_json]:
            if not osp.exists(json_file):
                continue
            coco = COCO(json_file)
            img_ids = coco.getImgIds()
            cat_ids = coco.getCatIds()
            anno_ids = coco.getAnnIds()

            catid2clsid = dict({catid: i for i, catid in enumerate(cat_ids)})
            clsid2catid = dict({i: catid for i, catid in enumerate(cat_ids)})
            cid2cname = dict({
                clsid: coco.loadCats(catid)[0]['name']
                for catid, clsid in catid2clsid.items()
            })

            # 由原train.json中的category生成新的category信息
            if json_file == train_files_json:
                cname2catid = dict({
                    coco.loadCats(catid)[0]['name']: clsid2catid[clsid]
                    for catid, clsid in catid2clsid.items()
                })
                new_cat_list = coco.loadCats(cat_ids)
            # 获取json中全部标注图片的名字
            for img_id in img_ids:
                img_anno = coco.loadImgs(img_id)[0]
                im_fname = img_anno['file_name']
                anno_set.add(im_fname)

                if json_file == train_files_json:
                    self.train_files.append(osp.join("JPEGImages", im_fname))
                elif json_file == val_files_json:
                    self.val_files.append(osp.join("JPEGImages", im_fname))
                elif json_file == test_files_json:
                    self.test_files.append(osp.join("JPEGImages", im_fname))
                # 获取每张图片的对应标注信息,并记录为npy格式
                anno_dict = read_coco_ann(img_id, coco, cid2cname, catid2clsid)
                img_path = osp.join("JPEGImages", im_fname)
                anno_path = osp.join("Annotations", im_fname)
                anno = replace_ext(anno_path, "npy")
                self.file_info[img_path] = anno

                # 生成label_info
                img_class = list(set(anno_dict["gt_class"]))
                for category_name in img_class:
                    if not category_name in self.label_info:
                        self.label_info[category_name] = [img_path]
                    else:
                        self.label_info[category_name].append(img_path)
                    if json_file == train_files_json:
                        if category_name in self.class_train_file_list:
                            self.class_train_file_list[category_name].append(
                                img_path)
                        else:
                            self.class_train_file_list[category_name] = list()
                            self.class_train_file_list[category_name].append(
                                img_path)
                    elif json_file == val_files_json:
                        if category_name in self.class_val_file_list:
                            self.class_val_file_list[category_name].append(
                                img_path)
                        else:
                            self.class_val_file_list[category_name] = list()
                            self.class_val_file_list[category_name].append(
                                img_path)
                    elif json_file == test_files_json:
                        if category_name in self.class_test_file_list:
                            self.class_test_file_list[category_name].append(
                                img_path)
                        else:
                            self.class_test_file_list[category_name] = list()
                            self.class_test_file_list[category_name].append(
                                img_path)

                # 生成新的图片信息
                new_img = img_anno
                new_img["id"] = img_id_index
                img_id_index += 1
                new_img_list.append(new_img)
                # 生成新的标注信息
                ins_anno_ids = coco.getAnnIds(imgIds=img_id, iscrowd=0)
                for ins_anno_id in ins_anno_ids:
                    anno = coco.loadAnns(ins_anno_id)[0]
                    new_anno = anno
                    new_anno["image_id"] = new_img["id"]
                    new_anno["id"] = anno_id_index
                    anno_id_index += 1
                    cat = coco.loadCats(anno["category_id"])[0]
                    new_anno_list.append(new_anno)

        if len(anno_set) > len(filename_set):
            sub_list = list(anno_set - filename_set)
            raise Exception("标注文件中{}等{}个信息无对应图片".format(sub_list[0],
                                                        len(sub_list)))

        for label in sorted(self.label_info.keys()):
            self.labels.append(label)

        self.annotation_dict = {
            "images": new_img_list,
            "categories": new_cat_list,
            "annotations": new_anno_list
        }

        # 若原数据集已切分，无annotations.json文件
        if not osp.exists(osp.join(self.path, "annotations.json")):
            json_file = open(osp.join(self.path, "annotations.json"), 'w+')
            json.dump(self.annotation_dict, json_file, cls=MyEncoder)
            json_file.close()

        # 生成每个图片对应的标注信息npy文件
        coco = COCO(osp.join(self.path, "annotations.json"))
        npy_path = osp.join(self.path, "Annotations")
        get_npy_from_coco_json(coco, npy_path, self.file_info)

        self.dump_statis_info()

    def split(self, val_split, test_split):
        all_files = list(self.file_info.keys())
        val_num = int(len(all_files) * val_split)
        test_num = int(len(all_files) * test_split)
        train_num = len(all_files) - val_num - test_num
        assert train_num > 0, "训练集样本数量需大于0"
        assert val_num > 0, "验证集样本数量需大于0"
        self.train_files = list()
        self.val_files = list()
        self.test_files = list()

        coco = COCO(osp.join(self.path, 'annotations.json'))
        img_ids = coco.getImgIds()
        cat_ids = coco.getCatIds()
        anno_ids = coco.getAnnIds()
        random.shuffle(img_ids)

        train_files_ids = img_ids[:train_num]
        val_files_ids = img_ids[train_num:train_num + val_num]
        test_files_ids = img_ids[train_num + val_num:]

        for img_id_list in [train_files_ids, val_files_ids, test_files_ids]:
            img_anno_ids = coco.getAnnIds(imgIds=img_id_list, iscrowd=0)
            imgs = coco.loadImgs(img_id_list)
            instances = coco.loadAnns(img_anno_ids)
            categories = coco.loadCats(cat_ids)
            img_dict = {
                "annotations": instances,
                "images": imgs,
                "categories": categories
            }

            if img_id_list == train_files_ids:
                for img in imgs:
                    self.train_files.append(
                        osp.join("JPEGImages", img["file_name"]))
                json_file = open(osp.join(self.path, 'train.json'), 'w+')
                json.dump(img_dict, json_file, cls=MyEncoder)
            elif img_id_list == val_files_ids:
                for img in imgs:
                    self.val_files.append(
                        osp.join("JPEGImages", img["file_name"]))
                json_file = open(osp.join(self.path, 'val.json'), 'w+')
                json.dump(img_dict, json_file, cls=MyEncoder)
            elif img_id_list == test_files_ids:
                for img in imgs:
                    self.test_files.append(
                        osp.join("JPEGImages", img["file_name"]))
                json_file = open(osp.join(self.path, 'test.json'), 'w+')
                json.dump(img_dict, json_file, cls=MyEncoder)

            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)

        self.dump_statis_info()