LMDBDataset is not in the dataset registry

[huiyang865]

你好，我在mmclassification v0.6.0版本基础上写了一个LMDB数据读取方式，但是当我创建一个LMDBDataset类的时候，代码提示错误：

• KeyError: 'LMDBDataset is not in the dataset registry'

如何注册一个数据类别？

我的LMDBDataset类别已经添加了@DATASETS.register_module()注册代码。

import mmcv
import lmdb
import numpy as np

from .builder import DATASETS
from .base_dataset import BaseDataset


@DATASETS.register_module()
class LMDBDataset(BaseDataset):

    def read_txt(self):
        data_infos = []
        with open(self.ann_file) as f:
            samples = [x.strip().split(' ') for x in f.readlines()]
            for filename, gt_label in samples:
                info = {'img_prefix': self.data_prefix}
                info['img_info'] = {'filename': filename}
                info['gt_label'] = np.array(gt_label, dtype=np.int64)
                data_infos.append(info)
            return data_infos

    def read_lmdb(self):
        data_infos = []
        env = lmdb.open(self.ann_file)
        txn = env.begin()
        class_num = int(txn.get('class_num'.encode()).decode())
        img_idx = 0
        for item_class in range(class_num):
            img_num = int(txn.get(f'class#{item_class}'.encode()).decode())
            for item_img in range(img_num):
                img_key = '###'.join(map(str, [img_idx, item_class])).encode()
                data_infos.append({'img_info': {'img_key': img_key},
                                   'gt_label': np.array(item_class, dtype=np.int64)})
        return data_infos

    def load_annotations(self):
        assert isinstance(self.ann_file, str)

        if self.ann_file.endswith('lmdb'):
            pass
        else:
            return self.read_txt()

感谢给予帮助。

[huiyang865]

需要在__init__.py文件中进一步添加相关类。已解决。