重写yolo数据加载模块

原因

 

        yolo原本的Dataload可阅读性太差，可拓展性也很差。另一个为了熟悉源代码。

class DataFolder(VisionDataset):
    def __init__(
            self,
            root: str,
            img_size=640,
            batch_size=32,
            augment=True,
            hyp=None,
            transform: Optional[Callable] = None,
            target_transform: Optional[Callable] = None
    ) -> None:
        super().__init__(root, transform=transform, target_transform=target_transform)
        self.sample = {}
        self.findLabelAndImg(root, self.sample)
        # 创建索引
        self.indices = range(len(self.sample['filePath']))
        self.img_size = img_size
        self.batch_size = batch_size
        # 超参数
        self.hyp = hyp
        self.augment = augment
        # 增强
        self.mosaic = True
        self.mixup = True
        self.albumentations = Albumentations(size=img_size) if augment else None

    def findLabelAndImg(self, root, sample):
        try:
            f = []  # image files
            p = Path(root)
            if p.is_dir():  # dir
                f += glob.glob(str(p / '**' / '*.*'), recursive=True)
            im_files = sorted(x.replace('/', os.sep) for x in f if x.split('.')[-1].lower() in IMG_FORMATS)
            assert im_files, f'No images found'
        except Exception as e:
            raise Exception(f'Error loading data from {root}: {e}') from e
        sample['filePath'] = f
        sample['labelPath'] = img2label_paths(im_files)


    def _load_image(self, id: int):
        f = self.sample['filePath'][id]
        im = cv2.imread(f)
        return im, im.shape[:2]  # BGR

    def _