# Ultralytics YOLO 🚀, AGPL-3.0 license

import contextlib
import json
from collections import defaultdict
from itertools import repeat
from multiprocessing.pool import ThreadPool
from pathlib import Path

import cv2
import numpy as np
import torch
from PIL import Image
from torch.utils.data import ConcatDataset

from ultralytics.utils import LOCAL_RANK, NUM_THREADS, TQDM, colorstr
from ultralytics.utils.ops import resample_segments

from .augment import (
    Compose,
    Format,
    Instances,
    LetterBox,
    RandomLoadText,
    classify_augmentations,
    classify_transforms,
    v8_transforms,
)
from .base import BaseDataset
from .utils import (
    HELP_URL,
    LOGGER,
    get_hash,
    img2label_paths,
    load_dataset_cache_file,
    save_dataset_cache_file,
    verify_image,
    verify_image_label,
)

# Ultralytics dataset *.cache version, >= 1.0.0 for YOLOv8
DATASET_CACHE_VERSION = "1.0.3"


class YOLODataset(BaseDataset):
    """
    Dataset class for loading object detection and/or segmentation labels in YOLO format.

    Args:
        data (dict, optional): A dataset YAML dictionary. Defaults to None.
        task (str): An explicit arg to point current task, Defaults to 'detect'.

    Returns:
        (torch.utils.data.Dataset): A PyTorch dataset object that can be used for training an object detection model.
    """

    def __init__(self, *args, data=None, task="detect", **kwargs):
        """Initializes the YOLODataset with optional configurations for segments and keypoints."""
        self.use_segments = task == "segment"
        self.use_keypoints = task == "pose"
        self.use_obb = task == "obb"
        self.data = data
        assert not (self.use_segments and self.use_keypoints), "Can not use both segments and keypoints."
        super().__init__(*args, **kwargs)

    def cache_labels(self, path=Path("./labels.cache")):
        """
        Cache dataset labels, check images and read shapes.

        Args:
            path (Path): Path where to save the cache file. Default is Path('./labels.cache').

        Returns:
            (dict): labels.
        """
        x = {"labels": []}
        nm, nf, ne, nc, msgs = 0, 0, 0, 0, []  # number missing, found, empty, corrupt, messages
        desc = f"{self.prefix}Scanning {path.parent / path.stem}..."
        total = len(self.im_files)
        nkpt, ndim = self.data.get("kpt_shape", (0, 0))
        if self.use_keypoints and (nkpt <= 0 or ndim not in {2, 3}):
            raise ValueError(
                "'kpt_shape' in data.yaml missing or incorrect. Should be a list with [number of "
                "keypoints, number of dims (2 for x,y or 3 for x,y,visible)], i.e. 'kpt_shape: [17, 3]'"
            )
        with ThreadPool(NUM_THREADS) as pool:
            results = pool.imap(
                func=verify_image_label,
                iterable=zip(
                    self.im_files,
                    self.label_files,
                    repeat(self.prefix),
                    repeat(self.use_keypoints),
                    repeat(len(self.data["names"])),
                    repeat(nkpt),
                    repeat(ndim),
                ),
            )
            pbar = TQDM(results, desc=desc, total=total)
            for im_file, lb, shape, segments, keypoint, nm_f, nf_f, ne_f, nc_f, msg in pbar:
                nm += nm_f
                nf += nf_f
                ne += ne_f
                nc += nc_f
                if im_file:
                    x["labels"].append(
                        {
                            "im_file": im_file,
                            "shape": shape,
                            "cls": lb[:, 0:1],  # n, 1
                            "bboxes": lb[:, 1:],  # n, 4
                            "segments": segments,
                            "keypoints": keypoint,
                            "normalized": True,
                            "bbox_format": "xywh",
                        }
                    )
                if msg:
                    msgs.append(msg)
                pbar.desc = f"{desc} {nf} images, {nm + ne} backgrounds, {nc} corrupt"
            pbar.close()

        if msgs:
            LOGGER.info("\n".join(msgs))
        if nf == 0:
            LOGGER.warning(f"{self.prefix}WARNING ⚠️ No labels found in {path}. {HELP_URL}")
        x["hash"] = get_hash(self.label_files + self.im_files)
        x["results"] = nf, nm, ne, nc, len(self.im_files)
        x["msgs"] = msgs  # warnings
        save_dataset_cache_file(self.prefix, path, x, DATASET_CACHE_VERSION)
        return x

    def get_labels(self):
        """Returns dictionary of labels for YOLO training."""
        self.label_files = img2label_paths(self.im_files)
        cache_path = Path(self.label_files[0]).parent.with_suffix(".cache")
        try:
            cache, exists = load_dataset_cache_file(cache_path), True  # attempt to load a *.cache file
            assert cache["version"] == DATASET_CACHE_VERSION  # matches current version
            assert cache["hash"] == get_hash(self.label_files + self.im_files)  # identical hash
        except (FileNotFoundError, AssertionError, AttributeError):
            cache, exists = self.cache_labels(cache_path), False  # run cache ops

        # Display cache
        nf, nm, ne, nc, n = cache.pop("results")  # found, missing, empty, corrupt, total
        if exists and LOCAL_RANK in {-1, 0}:
            d = f"Scanning {cache_path}... {nf} images, {nm + ne} backgrounds, {nc} corrupt"
            TQDM(None, desc=self.prefix + d, total=n, initial=n)  # display results
            if cache["msgs"]:
                LOGGER.info("\n".join(cache["msgs"]))  # display warnings

        # Read cache
        [cache.pop(k) for k in ("hash", "version", "msgs")]  # remove items
        labels = cache["labels"]
        if not labels:
            LOGGER.warning(f"WARNING ⚠️ No images found in {cache_path}, training may not work correctly. {HELP_URL}")
        self.im_files = [lb["im_file"] for lb in labels]  # update im_files

        # Check if the dataset is all boxes or all segments
        lengths = ((len(lb["cls"]), len(lb["bboxes"]), len(lb["segments"])) for lb in labels)
        len_cls, len_boxes, len_segments = (sum(x) for x in zip(*lengths))
        if len_segments and len_boxes != len_segments:
            LOGGER.warning(
                f"WARNING ⚠️ Box and segment counts should be equal, but got len(segments) = {len_segments}, "
                f"len(boxes) = {len_boxes}. To resolve this only boxes will be used and all segments will be removed. "
                "To avoid this please supply either a detect or segment dataset, not a detect-segment mixed dataset."
            )
            for lb in labels:
                lb["segments"] = []
        if len_cls == 0:
            LOGGER.warning(f"WARNING ⚠️ No labels found in {cache_path}, training may not work correctly. {HELP_URL}")
        return labels

    def build_transforms(self, hyp=None):
        """Builds and appends transforms to the list."""
        if self.augment:
            hyp.mosaic = hyp.mosaic if self.augment and not self.rect else 0.0
            hyp.mixup = hyp.mixup if self.augment and not self.rect else 0.0
            transforms = v8_transforms(self, self.imgsz, hyp)
        else:
            transforms = Compose([LetterBox(new_shape=(self.imgsz, self.imgsz), scaleup=False)])
        transforms.append(
            Format(
                bbox_format="xywh",
                normalize=True,
                return_mask=self.use_segments,
                return_keypoint=self.use_keypoints,
                return_obb=self.use_obb,
                batch_idx=True,
                mask_ratio=hyp.mask_ratio,
                mask_overlap=hyp.overlap_mask,
                bgr=1.0 # hyp.bgr if self.augment else 0.0,  # only affect training.
            )
        )
        return transforms

    def close_mosaic(self, hyp):
        """Sets mosaic, copy_paste and mixup options to 0.0 and builds transformations."""
        hyp.mosaic = 0.0  # set mosaic ratio=0.0
        hyp.copy_paste = 0.0  # keep the same behavior as previous v8 close-mosaic
        hyp.mixup = 0.0  # keep the same behavior as previous v8 close-mosaic
        self.transforms = self.build_transforms(hyp)

    def update_labels_info(self, label):
        """
        Custom your label format here.

        Note:
            cls is not with bboxes now, classification and semantic segmentation need an independent cls label
            Can also support classification and semantic segmentation by adding or removing dict keys there.
        """
        bboxes = label.pop("bboxes")
        segments = label.pop("segments", [])
        keypoints = label.pop("keypoints", None)
        bbox_format = label.pop("bbox_format")
        normalized = label.pop("normalized")

        # NOTE: do NOT resample oriented boxes
        segment_resamples = 100 if self.use_obb else 1000
        if len(segments) > 0:
            # list[np.array(1000, 2)] * num_samples
            # (N, 1000, 2)
            segments = np.stack(resample_segments(segments, n=segment_resamples), axis=0)
        else:
            segments = np.zeros((0, segment_resamples, 2), dtype=np.float32)
        label["instances"] = Instances(bboxes, segments, keypoints, bbox_format=bbox_format, normalized=normalized)
        return label

    @staticmethod
    def collate_fn(batch):
        """Collates data samples into batches."""
        new_batch = {}
        keys = batch[0].keys()
        values = list(zip(*[list(b.values()) for b in batch]))
        for i, k in enumerate(keys):
            value = values[i]
            if k == "img":
                value = torch.stack(value, 0)
            if k in {"masks", "keypoints", "bboxes", "cls", "bboxes_map", "segments", "obb"}:
                value = [v for v in value if v.numel() > 0] #zsk
                if len(value) == 0:
                    value = torch.empty(0)
                else:
                    value = torch.cat(value, 0)
            new_batch[k] = value
        new_batch["batch_idx"] = list(new_batch["batch_idx"])
        for i in range(len(new_batch["batch_idx"])):
            new_batch["batch_idx"][i] += i  # add target image index for build_targets()
        new_batch["batch_idx"] = torch.cat(new_batch["batch_idx"], 0)
        return new_batch

import random
import math
from scipy.stats import norm
class YOLODetrackDataset(YOLODataset):
    """
    Dataset class for loading object detection and/or segmentation labels in YOLO format.

    Args:
        data (dict, optional): A dataset YAML dictionary. Defaults to None.
        task (str): An explicit arg to point current task, Defaults to 'detect'.

    Returns:
        (torch.utils.data.Dataset): A PyTorch dataset object that can be used for training an object detection model.
    """
    def __init__(self, *args, img_path, data=None, task="detect", **kwargs):
        """Initializes a dataset object for object detection tasks with optional specifications."""
        self.N = 1
        if isinstance(img_path, (list, tuple)):
            self.N = len(img_path)
        if self.N <= 1:
            super().__init__(*args, img_path, data=data, task=task, **kwargs)
        else:
            self.datasets = [
                YOLODataset(*args, img_path=path, data=data, task=task, **kwargs)
                for path in img_path[1:]
            ]
            super().__init__(*args, img_path=img_path[0], data=data, task=task, **kwargs)

    def get_image_and_label(self, index):
        """Get and return label information from the dataset."""
        label = deepcopy(self.labels[index])  # requires deepcopy() https://github.com/ultralytics/ultralytics/pull/1948
        label.pop("shape", None)  # shape is for rect, remove it
        label["img"], label["ori_shape"], label["resized_shape"] = self.load_image(index)
        label["ratio_pad"] = (
            label["resized_shape"][0] / label["ori_shape"][0],
            label["resized_shape"][1] / label["ori_shape"][1],
        )  # for evaluation
        if self.rect:
            label["rect_shape"] = self.batch_shapes[self.batch[index]]

        if self.N == 2: # 单通道mask，有且仅有一个
            height, width, channels = label["img"].shape
            mask = np.zeros((height, width), dtype=label["img"].dtype)
            if self.augment is not True: # 验证时保持一致
                mask, _, _ = self.datasets[0].load_image(index)
            elif random.random() < 0.6:
                mask, _, _ = self.datasets[0].load_image(index)
            label["img"] = np.concatenate((label["img"], mask[:,:,np.newaxis]), axis=-1)
            return self.update_labels_info(label)
        elif self.N > 2:
            # 转换为单通道以节约内存
            label["img"] = np.mean(label["img"], axis=2, keepdims=True).astype(np.uint8)
            for idx, dataset in enumerate(self.datasets):
                labeli = deepcopy(dataset.labels[index])
                labeli["img"], _, _ = dataset.load_image(index)
                labeli["img"] = np.mean(labeli["img"], axis=2, keepdims=True).astype(np.uint8)
                
                label["img"] = np.concatenate((labeli["img"], label["img"]), axis=-1)
                # label["im_file"] = ([label["im_file"]] if isinstance(label["im_file"], str) else label["im_file"]) + [labeli["im_file"]]
            return self.update_labels_info(label)
        
        height, width, channels = label["img"].shape
        mask = np.zeros((height, width, 1), dtype=label["img"].dtype)
        # if self.augment is not True:
        #     # 验证时不添加噪声容易评估
        #     # 将标签嵌入mask，标签格式为归一化x,y,w,h，图像一般为同比例缩放h,w
        #     for i in range(label["bboxes"].shape[0]):
        #         x_center, y_center, box_width, box_height = label["bboxes"][i, ...]
        #         # 反归一化坐标和尺寸，以像素值表示
        #         x_center = int(x_center * width)
        #         y_center = int(y_center * height)
        #         box_width = int(box_width * width)
        #         box_height = int(box_height * height) 
        #         # 计算左上角和右下角坐标
        #         x1 = int(x_center - box_width / 2)
        #         y1 = int(y_center - box_height / 2)
        #         x2 = int(x_center + box_width / 2)
        #         y2 = int(y_center + box_height / 2)
        #         # 将掩模中目标区域置为 1
        #         mask[y1:y2, x1:x2] = 1
        # else:
        # 添加噪声，训练时添加噪声
        for i in range(label["bboxes"].shape[0]):
            x_center, y_center, box_width, box_height = label["bboxes"][i, ...]
            x_center = int(x_center * width)
            y_center = int(y_center * height)
            box_width = int(box_width * width)
            box_height = int(box_height * height)
            if box_width==0 or box_height==0:
                continue
            # # 随机缩放和偏移参数
            # scale_factor = 8 # 0.125 0.25 0.5 1 2 4 8 16 32
            # offset = 4 # random.uniform(0, abs(scale_factor - 1)) # 0-0.5
            # # 标注框位置随机偏移
            # rad = random.uniform(0, 2*math.pi)
            # distance = offset * math.sqrt(box_width**2 + box_height**2) / 2.0 #整框半对角线
            # x_center += int(distance * math.cos(rad))
            # y_center += int(distance * math.sin(rad))
            # # 标注框大小随机缩放
            # box_width = int(box_width * scale_factor)
            # box_height = int(box_height * scale_factor)
            # 计算左上角和右下角坐标
            x1 = max(int(x_center - box_width / 2), 0)
            y1 = max(int(y_center - box_height / 2), 0)
            x2 = min(int(x_center + box_width / 2), width)
            y2 = min(int(y_center + box_height / 2), height)
            # 掩模存在概率 与假掩膜概率正相关
            if random.random() < 0.5:
                # mask[y1:y2, x1:x2] = 255
                mask_patch = self.create_gaussian_patch(box_height, box_width)
                mask[y1:y2, x1:x2] = np.maximum(mask[y1:y2, x1:x2], mask_patch[:y2-y1, :x2-x1, np.newaxis]) # 使用最大值融合
        # 随机增加假的目标掩膜
        fake_p = 0.5
        if random.random() < fake_p:
            # rng = math.ceil(fake_p*10*0.6)
            fake_num = random.randint(0, 20) # 5
            fake_num = fake_num if fake_num>0 else 0
        else: 
            fake_num = 0
        for i in range(fake_num):
            fake_x_center = random.randint(0, width)
            fake_y_center = random.randint(0, height)
            fake_box_width = random.randint(width // 100, width // 40) #40 20
            fake_box_height = random.randint(height // 100, height // 40) #40 20
            if fake_box_width==0 or fake_box_height==0:
                continue
            fx1 = max(int(fake_x_center - fake_box_width / 2), 0)
            fy1 = max(int(fake_y_center - fake_box_height / 2), 0)
            fx2 = min(int(fake_x_center + fake_box_width / 2), width)
            fy2 = min(int(fake_y_center + fake_box_height / 2), height)
            # mask[fy1:fy2, fx1:fx2] = 255
            mask_patch = self.create_gaussian_patch(fake_box_height, fake_box_width)
            mask[fy1:fy2, fx1:fx2] = np.maximum(mask[fy1:fy2, fx1:fx2], mask_patch[:fy2-fy1, :fx2-fx1, np.newaxis]) # 使用最大值融合
        label["img"] = np.concatenate((label["img"], mask), axis=-1)
        
        return self.update_labels_info(label)
    
    def build_transforms(self, hyp=None):
        """Builds and appends transforms to the list."""
        if self.augment:
            hyp.mosaic = hyp.mosaic if self.augment and not self.rect else 0.0
            hyp.mixup = hyp.mixup if self.augment and not self.rect else 0.0
            transforms = v8_transforms(self, self.imgsz, hyp)
        else:
            transforms = Compose([LetterBox(new_shape=(self.imgsz, self.imgsz), scaleup=False)])
        transforms.append(
            Format(
                bbox_format="xywh",
                normalize=True,
                return_mask=self.use_segments,
                return_keypoint=self.use_keypoints,
                return_obb=self.use_obb,
                batch_idx=True,
                mask_ratio=hyp.mask_ratio,
                mask_overlap=hyp.overlap_mask,
                bgr=1, # bgr=0 if self.N<=2 else 1,  # 1 no bgr2rgb
            )
        )
        return transforms
    
    def create_gaussian_patch(self, height, width):
        if height==0 or width==0:
            print("size is 0")
        # 计算高斯分布的标准差（控制分布宽度）
        sigma_y, sigma_x = height / 6, width / 6  # 约99.7%的值将在[-3sigma, 3sigma]范围内
        
        # 创建坐标网格
        y = np.arange(0, height)
        x = np.arange(0, width)
        x, y = np.meshgrid(x, y)
        
        # 计算高斯分布值
        center_y, center_x = height // 2, width // 2  # 局部矩阵中心
        gaussian_patch = norm.pdf(np.sqrt((x - center_x)**2 / sigma_x**2 + (y - center_y)**2 / sigma_y**2))
        
        # 归一化到0-1范围
        gaussian_patch = gaussian_patch - np.min(gaussian_patch)
        gaussian_patch = gaussian_patch / np.max(gaussian_patch)

         # 归一化到0-255范围
        gaussian_patch = (gaussian_patch * 255).astype(np.uint8)
        
        return gaussian_patch

class YOLODetrackDataset1(YOLODataset):
    """
    Dataset class for loading object detection and/or segmentation labels in YOLO format.

    Args:
        data (dict, optional): A dataset YAML dictionary. Defaults to None.
        task (str): An explicit arg to point current task, Defaults to 'detect'.

    Returns:
        (torch.utils.data.Dataset): A PyTorch dataset object that can be used for training an object detection model.
    """
    def __init__(self, *args, data=None, task="detect", **kwargs):
        """Initializes a dataset object for object detection tasks with optional specifications."""
        super().__init__(*args, data=data, task=task, **kwargs)

    def get_image_and_label(self, index):
        """Get and return label information from the dataset."""
        label = deepcopy(self.labels[index])  # requires deepcopy() https://github.com/ultralytics/ultralytics/pull/1948
        label.pop("shape", None)  # shape is for rect, remove it
        label["img"], label["ori_shape"], label["resized_shape"] = self.load_image(index)
        label["ratio_pad"] = (
            label["resized_shape"][0] / label["ori_shape"][0],
            label["resized_shape"][1] / label["ori_shape"][1],
        )  # for evaluation
        if self.rect:
            label["rect_shape"] = self.batch_shapes[self.batch[index]]
        
        height, width, channels = label["img"].shape
        mask = np.zeros((height, width, 1), dtype=label["img"].dtype)
        for i in range(label["bboxes"].shape[0]):
            x_center, y_center, box_width, box_height = label["bboxes"][i, ...]
            x_center = int(x_center * width)
            y_center = int(y_center * height)
            box_width = int(box_width * width)
            box_height = int(box_height * height)
            if box_width==0 or box_height==0:
                continue
            # 随机缩放和偏移参数
            scale_factor = 1 # 0.125 0.25 0.5 1 2 4 8 16 32
            offset = random.uniform(0, scale_factor*0.25) # 0-2-8
            # 标注框位置随机偏移
            rad = random.uniform(0, 2*math.pi)
            distance = offset * math.sqrt(box_width**2 + box_height**2) / 2.0 #整框半对角线
            x_center += int(distance * math.cos(rad))
            y_center += int(distance * math.sin(rad))
            # 标注框大小缩放
            box_width = int(box_width * scale_factor)
            box_height = int(box_height * scale_factor)
            # 计算左上角和右下角坐标
            x1 = max(int(x_center - box_width / 2), 0)
            y1 = max(int(y_center - box_height / 2), 0)
            x2 = min(int(x_center + box_width / 2), width)
            y2 = min(int(y_center + box_height / 2), height)
            # 掩模存在概率 与假掩膜概率正相关
            if random.random() < 0.6:
                mask[y1:y2, x1:x2] = 255
                # mask_patch = self.create_gaussian_patch(box_height, box_width)
                # mask[y1:y2, x1:x2] = np.maximum(mask[y1:y2, x1:x2], mask_patch[:y2-y1, :x2-x1, np.newaxis]) # 使用最大值融合
        # 随机增加假的目标掩膜
        fake_p = 0.6
        if random.random() < fake_p:
            # rng = math.ceil(fake_p*10*0.6)
            fake_num = random.randint(1, 2) # 5
            fake_num = fake_num if fake_num>0 else 0
        else: 
            fake_num = 0
        for i in range(fake_num):
            fake_x_center = random.randint(0, width)
            fake_y_center = random.randint(0, height)
            fake_box_width = random.randint(width // 100, width // 40) #40 20
            fake_box_height = random.randint(height // 100, height // 40) #40 20
            if fake_box_width==0 or fake_box_height==0:
                continue
            fx1 = max(int(fake_x_center - fake_box_width / 2), 0)
            fy1 = max(int(fake_y_center - fake_box_height / 2), 0)
            fx2 = min(int(fake_x_center + fake_box_width / 2), width)
            fy2 = min(int(fake_y_center + fake_box_height / 2), height)
            mask[fy1:fy2, fx1:fx2] = 255
            # mask_patch = self.create_gaussian_patch(fake_box_height, fake_box_width)
            # mask[fy1:fy2, fx1:fx2] = np.maximum(mask[fy1:fy2, fx1:fx2], mask_patch[:fy2-fy1, :fx2-fx1, np.newaxis]) # 使用最大值融合
        label["img"] = np.concatenate((label["img"], mask), axis=-1)
        
        return self.update_labels_info(label)
    
    def build_transforms(self, hyp=None):
        """Builds and appends transforms to the list."""
        if self.augment:
            hyp.mosaic = hyp.mosaic if self.augment and not self.rect else 0.0
            hyp.mixup = hyp.mixup if self.augment and not self.rect else 0.0
            transforms = v8_transforms(self, self.imgsz, hyp)
        else:
            transforms = Compose([LetterBox(new_shape=(self.imgsz, self.imgsz), scaleup=False)])
        transforms.append(
            Format(
                bbox_format="xywh",
                normalize=True,
                return_mask=self.use_segments,
                return_keypoint=self.use_keypoints,
                return_obb=self.use_obb,
                batch_idx=True,
                mask_ratio=hyp.mask_ratio,
                mask_overlap=hyp.overlap_mask,
                bgr=1,  # 1 no bgr2rgb
            )
        )
        return transforms
    
    def create_gaussian_patch(self, height, width):
        if height==0 or width==0:
            print("size is 0")
        # 计算高斯分布的标准差（控制分布宽度）
        sigma_y, sigma_x = height / 6, width / 6  # 约99.7%的值将在[-3sigma, 3sigma]范围内
        
        # 创建坐标网格
        y = np.arange(0, height)
        x = np.arange(0, width)
        x, y = np.meshgrid(x, y)
        
        # 计算高斯分布值
        center_y, center_x = height // 2, width // 2  # 局部矩阵中心
        gaussian_patch = norm.pdf(np.sqrt((x - center_x)**2 / sigma_x**2 + (y - center_y)**2 / sigma_y**2))
        
        # 归一化到0-1范围
        gaussian_patch = gaussian_patch - np.min(gaussian_patch)
        gaussian_patch = gaussian_patch / np.max(gaussian_patch)

         # 归一化到0-255范围
        gaussian_patch = (gaussian_patch * 255).astype(np.uint8)
        
        return gaussian_patch


import random
import math
from scipy.stats import norm
class YOLODemaskDataset(YOLODataset):
    """
    Dataset class for loading object detection and/or segmentation labels in YOLO format.

    Args:
        data (dict, optional): A dataset YAML dictionary. Defaults to None.
        task (str): An explicit arg to point current task, Defaults to 'detect'.

    Returns:
        (torch.utils.data.Dataset): A PyTorch dataset object that can be used for training an object detection model.
    """
    def __init__(self, *args, img_path, data=None, task="detect", **kwargs):
        """Initializes a dataset object for object detection tasks with optional specifications."""
        self.N = 1
        if isinstance(img_path, (list, tuple)):
            self.N = len(img_path)
        if self.N <= 1:
            super().__init__(*args, img_path, data=data, task=task, **kwargs)
        else:
            self.datasets = [
                YOLODataset(*args, img_path=path, data=data, task=task, **kwargs)
                for path in img_path[1:]
            ]
            super().__init__(*args, img_path=img_path[0], data=data, task=task, **kwargs)

    def get_image_and_label(self, index):
        """Get and return label information from the dataset."""
        label = deepcopy(self.labels[index])  # requires deepcopy() https://github.com/ultralytics/ultralytics/pull/1948
        label.pop("shape", None)  # shape is for rect, remove it
        label["img"], label["ori_shape"], label["resized_shape"] = self.load_image(index)
        label["ratio_pad"] = (
            label["resized_shape"][0] / label["ori_shape"][0],
            label["resized_shape"][1] / label["ori_shape"][1],
        )  # for evaluation
        if self.rect:
            label["rect_shape"] = self.batch_shapes[self.batch[index]]

        if self.N > 1:
            # 转换为单通道以节约内存
            label["img"] = np.mean(label["img"], axis=2, keepdims=True).astype(np.uint8)
            for idx, dataset in enumerate(self.datasets):
                labeli = deepcopy(dataset.labels[index])
                labeli["img"], _, _ = dataset.load_image(index)
                labeli["img"] = np.mean(labeli["img"], axis=2, keepdims=True).astype(np.uint8)
                
                label["img"] = np.concatenate((labeli["img"], label["img"]), axis=-1)
                # label["im_file"] = ([label["im_file"]] if isinstance(label["im_file"], str) else label["im_file"]) + [labeli["im_file"]]

        height, width, channels = label["img"].shape
        mask = np.zeros((height, width, 1), dtype=label["img"].dtype)
        for i in range(label["bboxes"].shape[0]):
            x_center, y_center, box_width, box_height = label["bboxes"][i, ...]
            x_center = int(x_center * width)
            y_center = int(y_center * height)
            box_width = int(box_width * width)
            box_height = int(box_height * height)
            if box_width==0 or box_height==0:
                continue
            # 计算左上角和右下角坐标
            x1 = max(int(x_center - box_width / 2), 0)
            y1 = max(int(y_center - box_height / 2), 0)
            x2 = min(int(x_center + box_width / 2), width)
            y2 = min(int(y_center + box_height / 2), height)
            # 掩模
            mask_patch = self.create_gaussian_patch(box_height, box_width)
            mask[y1:y2, x1:x2] = np.maximum(mask[y1:y2, x1:x2], mask_patch[:y2-y1, :x2-x1, np.newaxis]) # 使用最大值融合
        
        label["img"] = np.concatenate((label["img"], mask), axis=-1)
        
        return self.update_labels_info(label)
    
    def build_transforms(self, hyp=None):
        """Builds and appends transforms to the list."""
        if self.augment:
            hyp.mosaic = hyp.mosaic if self.augment and not self.rect else 0.0
            hyp.mixup = hyp.mixup if self.augment and not self.rect else 0.0
            transforms = v8_transforms(self, self.imgsz, hyp)
        else:
            transforms = Compose([LetterBox(new_shape=(self.imgsz, self.imgsz), scaleup=False)])
        transforms.append(
            Format(
                bbox_format="xywh",
                normalize=True,
                return_mask=self.use_segments,
                return_keypoint=self.use_keypoints,
                return_obb=self.use_obb,
                batch_idx=True,
                mask_ratio=hyp.mask_ratio,
                mask_overlap=hyp.overlap_mask,
                bgr=0 if self.N==1 else 1,  # 1 no bgr2rgb
                demask=True,
            )
        )
        return transforms
    
    def create_gaussian_patch(self, height, width):
        if height==0 or width==0:
            print("size is 0")

        # # 外接椭圆
        # a = int(np.sqrt(height**2 + width**2))
        # b = max(height, width)
        # height, width = (a, b) if height > width else (b, a)

        # 计算高斯分布的标准差（控制分布宽度）
        sigma_y, sigma_x = height / 6, width / 6  # 约99.7%的值将在[-3sigma, 3sigma]范围内
        
        # 创建坐标网格
        y = np.arange(0, height)
        x = np.arange(0, width)
        x, y = np.meshgrid(x, y)
        
        # 计算高斯分布值
        center_y, center_x = height // 2, width // 2  # 局部矩阵中心
        gaussian_patch = norm.pdf(np.sqrt((x - center_x)**2 / sigma_x**2 + (y - center_y)**2 / sigma_y**2))
        
        # 归一化到0-1范围
        gaussian_patch = gaussian_patch - np.min(gaussian_patch)
        gaussian_patch = gaussian_patch / np.max(gaussian_patch)

         # 归一化到0-255范围
        gaussian_patch = np.round(gaussian_patch * 255).astype(np.uint8) #四舍五入
        
        return gaussian_patch


from copy import deepcopy
class YOLOMultiFrameDataset(YOLODataset):
    """
    Dataset class for loading object detection and/or segmentation labels in YOLO format.

    Args:
        data (dict, optional): A dataset YAML dictionary. Defaults to None.
        task (str): An explicit arg to point current task, Defaults to 'detect'.

    Returns:
        (torch.utils.data.Dataset): A PyTorch dataset object that can be used for training an object detection model.
    """
    def __init__(self, *args, img_path, data=None, task="detect", **kwargs):
        """Initializes the YOLODataset with optional configurations for segments and keypoints."""
        self.N = len(img_path)
        self.datasets = [
            YOLODataset(*args, img_path=path, data=data, task=task, **kwargs)
            for path in img_path[1:]
        ]
        super().__init__(*args, img_path=img_path[0], data=data, task=task, **kwargs)

    def get_image_and_label(self, index):
        """Get and return label information from the dataset."""
        label = deepcopy(self.labels[index])
        label.pop("shape", None)  # shape is for rect, remove it
        label["img"], label["ori_shape"], label["resized_shape"] = self.load_image(index)
        label["ratio_pad"] = (
            label["resized_shape"][0] / label["ori_shape"][0],
            label["resized_shape"][1] / label["ori_shape"][1],
        )  # for evaluation
        if self.rect:
            label["rect_shape"] = self.batch_shapes[self.batch[index]]

        label["bboxes_map"] = np.zeros((label["bboxes"].shape[0], 1))
        for idx, dataset in enumerate(self.datasets):
            labeli = deepcopy(dataset.labels[index])
            labeli["img"], _, _ = dataset.load_image(index)

            label["img"] = cv2.merge([label["img"], labeli["img"]])
            label["bboxes"] = np.concatenate((label["bboxes"], labeli["bboxes"]), axis=0)
            bboxes_mapi = np.ones((labeli["bboxes"].shape[0], 1)) * (idx+1)
            label["bboxes_map"] = np.concatenate((label["bboxes_map"], bboxes_mapi), axis=0)
            label["cls"] = np.concatenate((label["cls"], labeli["cls"]), axis=0)
            # label["im_file"] = ([label["im_file"]] if isinstance(label["im_file"], str) else label["im_file"]) + [labeli["im_file"]]

        return self.update_labels_info(label)
    
    def update_labels_info(self, label):
        """
        Custom your label format here.

        Note:
            cls is not with bboxes now, classification and semantic segmentation need an independent cls label
            Can also support classification and semantic segmentation by adding or removing dict keys there.
        """
        bboxes_map = label.pop("bboxes_map")
        bboxes = label.pop("bboxes")
        segments = label.pop("segments", [])
        keypoints = label.pop("keypoints", None)
        bbox_format = label.pop("bbox_format")
        normalized = label.pop("normalized")

        # NOTE: do NOT resample oriented boxes
        segment_resamples = 100 if self.use_obb else 1000
        if len(segments) > 0:
            # list[np.array(1000, 2)] * num_samples
            # (N, 1000, 2)
            segments = np.stack(resample_segments(segments, n=segment_resamples), axis=0)
        else:
            segments = np.zeros((0, segment_resamples, 2), dtype=np.float32)
        label["instances"] = Instances(bboxes, segments, keypoints, bbox_format=bbox_format, normalized=normalized, bboxes_map=bboxes_map)
        return label
    

class YOLOMultiModalDataset(YOLODataset):
    """
    Dataset class for loading object detection and/or segmentation labels in YOLO format.

    Args:
        data (dict, optional): A dataset YAML dictionary. Defaults to None.
        task (str): An explicit arg to point current task, Defaults to 'detect'.

    Returns:
        (torch.utils.data.Dataset): A PyTorch dataset object that can be used for training an object detection model.
    """

    def __init__(self, *args, data=None, task="detect", **kwargs):
        """Initializes a dataset object for object detection tasks with optional specifications."""
        super().__init__(*args, data=data, task=task, **kwargs)

    def update_labels_info(self, label):
        """Add texts information for multi modal model training."""
        labels = super().update_labels_info(label)
        # NOTE: some categories are concatenated with its synonyms by `/`.
        labels["texts"] = [v.split("/") for _, v in self.data["names"].items()]
        return labels

    def build_transforms(self, hyp=None):
        """Enhances data transformations with optional text augmentation for multi-modal training."""
        transforms = super().build_transforms(hyp)
        if self.augment:
            # NOTE: hard-coded the args for now.
            transforms.insert(-1, RandomLoadText(max_samples=min(self.data["nc"], 80), padding=True))
        return transforms


class GroundingDataset(YOLODataset):
    def __init__(self, *args, task="detect", json_file, **kwargs):
        """Initializes a GroundingDataset for object detection, loading annotations from a specified JSON file."""
        assert task == "detect", "`GroundingDataset` only support `detect` task for now!"
        self.json_file = json_file
        super().__init__(*args, task=task, data={}, **kwargs)

    def get_img_files(self, img_path):
        """The image files would be read in `get_labels` function, return empty list here."""
        return []

    def get_labels(self):
        """Loads annotations from a JSON file, filters, and normalizes bounding boxes for each image."""
        labels = []
        LOGGER.info("Loading annotation file...")
        with open(self.json_file, "r") as f:
            annotations = json.load(f)
        images = {f'{x["id"]:d}': x for x in annotations["images"]}
        imgToAnns = defaultdict(list)
        for ann in annotations["annotations"]:
            imgToAnns[ann["image_id"]].append(ann)
        for img_id, anns in TQDM(imgToAnns.items(), desc=f"Reading annotations {self.json_file}"):
            img = images[f"{img_id:d}"]
            h, w, f = img["height"], img["width"], img["file_name"]
            im_file = Path(self.img_path) / f
            if not im_file.exists():
                continue
            self.im_files.append(str(im_file))
            bboxes = []
            cat2id = {}
            texts = []
            for ann in anns:
                if ann["iscrowd"]:
                    continue
                box = np.array(ann["bbox"], dtype=np.float32)
                box[:2] += box[2:] / 2
                box[[0, 2]] /= float(w)
                box[[1, 3]] /= float(h)
                if box[2] <= 0 or box[3] <= 0:
                    continue

                cat_name = " ".join([img["caption"][t[0] : t[1]] for t in ann["tokens_positive"]])
                if cat_name not in cat2id:
                    cat2id[cat_name] = len(cat2id)
                    texts.append([cat_name])
                cls = cat2id[cat_name]  # class
                box = [cls] + box.tolist()
                if box not in bboxes:
                    bboxes.append(box)
            lb = np.array(bboxes, dtype=np.float32) if len(bboxes) else np.zeros((0, 5), dtype=np.float32)
            labels.append(
                {
                    "im_file": im_file,
                    "shape": (h, w),
                    "cls": lb[:, 0:1],  # n, 1
                    "bboxes": lb[:, 1:],  # n, 4
                    "normalized": True,
                    "bbox_format": "xywh",
                    "texts": texts,
                }
            )
        return labels

    def build_transforms(self, hyp=None):
        """Configures augmentations for training with optional text loading; `hyp` adjusts augmentation intensity."""
        transforms = super().build_transforms(hyp)
        if self.augment:
            # NOTE: hard-coded the args for now.
            transforms.insert(-1, RandomLoadText(max_samples=80, padding=True))
        return transforms


class YOLOConcatDataset(ConcatDataset):
    """
    Dataset as a concatenation of multiple datasets.

    This class is useful to assemble different existing datasets.
    """

    @staticmethod
    def collate_fn(batch):
        """Collates data samples into batches."""
        return YOLODataset.collate_fn(batch)


# TODO: support semantic segmentation
class SemanticDataset(BaseDataset):
    """
    Semantic Segmentation Dataset.

    This class is responsible for handling datasets used for semantic segmentation tasks. It inherits functionalities
    from the BaseDataset class.

    Note:
        This class is currently a placeholder and needs to be populated with methods and attributes for supporting
        semantic segmentation tasks.
    """

    def __init__(self):
        """Initialize a SemanticDataset object."""
        super().__init__()


class ClassificationDataset:
    """
    Extends torchvision ImageFolder to support YOLO classification tasks, offering functionalities like image
    augmentation, caching, and verification. It's designed to efficiently handle large datasets for training deep
    learning models, with optional image transformations and caching mechanisms to speed up training.

    This class allows for augmentations using both torchvision and Albumentations libraries, and supports caching images
    in RAM or on disk to reduce IO overhead during training. Additionally, it implements a robust verification process
    to ensure data integrity and consistency.

    Attributes:
        cache_ram (bool): Indicates if caching in RAM is enabled.
        cache_disk (bool): Indicates if caching on disk is enabled.
        samples (list): A list of tuples, each containing the path to an image, its class index, path to its .npy cache
                        file (if caching on disk), and optionally the loaded image array (if caching in RAM).
        torch_transforms (callable): PyTorch transforms to be applied to the images.
    """

    def __init__(self, root, args, augment=False, prefix=""):
        """
        Initialize YOLO object with root, image size, augmentations, and cache settings.

        Args:
            root (str): Path to the dataset directory where images are stored in a class-specific folder structure.
            args (Namespace): Configuration containing dataset-related settings such as image size, augmentation
                parameters, and cache settings. It includes attributes like `imgsz` (image size), `fraction` (fraction
                of data to use), `scale`, `fliplr`, `flipud`, `cache` (disk or RAM caching for faster training),
                `auto_augment`, `hsv_h`, `hsv_s`, `hsv_v`, and `crop_fraction`.
            augment (bool, optional): Whether to apply augmentations to the dataset. Default is False.
            prefix (str, optional): Prefix for logging and cache filenames, aiding in dataset identification and
                debugging. Default is an empty string.
        """
        import torchvision  # scope for faster 'import ultralytics'

        # Base class assigned as attribute rather than used as base class to allow for scoping slow torchvision import
        self.base = torchvision.datasets.ImageFolder(root=root)
        self.samples = self.base.samples
        self.root = self.base.root

        # Initialize attributes
        if augment and args.fraction < 1.0:  # reduce training fraction
            self.samples = self.samples[: round(len(self.samples) * args.fraction)]
        self.prefix = colorstr(f"{prefix}: ") if prefix else ""
        self.cache_ram = args.cache is True or str(args.cache).lower() == "ram"  # cache images into RAM
        self.cache_disk = str(args.cache).lower() == "disk"  # cache images on hard drive as uncompressed *.npy files
        self.samples = self.verify_images()  # filter out bad images
        self.samples = [list(x) + [Path(x[0]).with_suffix(".npy"), None] for x in self.samples]  # file, index, npy, im
        scale = (1.0 - args.scale, 1.0)  # (0.08, 1.0)
        self.torch_transforms = (
            classify_augmentations(
                size=args.imgsz,
                scale=scale,
                hflip=args.fliplr,
                vflip=args.flipud,
                erasing=args.erasing,
                auto_augment=args.auto_augment,
                hsv_h=args.hsv_h,
                hsv_s=args.hsv_s,
                hsv_v=args.hsv_v,
            )
            if augment
            else classify_transforms(size=args.imgsz, crop_fraction=args.crop_fraction)
        )

    def __getitem__(self, i):
        """Returns subset of data and targets corresponding to given indices."""
        f, j, fn, im = self.samples[i]  # filename, index, filename.with_suffix('.npy'), image
        if self.cache_ram:
            if im is None:  # Warning: two separate if statements required here, do not combine this with previous line
                im = self.samples[i][3] = cv2.imread(f)
        elif self.cache_disk:
            if not fn.exists():  # load npy
                np.save(fn.as_posix(), cv2.imread(f), allow_pickle=False)
            im = np.load(fn)
        else:  # read image
            im = cv2.imread(f)  # BGR
        # Convert NumPy array to PIL image
        im = Image.fromarray(cv2.cvtColor(im, cv2.COLOR_BGR2RGB))
        sample = self.torch_transforms(im)
        return {"img": sample, "cls": j}

    def __len__(self) -> int:
        """Return the total number of samples in the dataset."""
        return len(self.samples)

    def verify_images(self):
        """Verify all images in dataset."""
        desc = f"{self.prefix}Scanning {self.root}..."
        path = Path(self.root).with_suffix(".cache")  # *.cache file path

        with contextlib.suppress(FileNotFoundError, AssertionError, AttributeError):
            cache = load_dataset_cache_file(path)  # attempt to load a *.cache file
            assert cache["version"] == DATASET_CACHE_VERSION  # matches current version
            assert cache["hash"] == get_hash([x[0] for x in self.samples])  # identical hash
            nf, nc, n, samples = cache.pop("results")  # found, missing, empty, corrupt, total
            if LOCAL_RANK in {-1, 0}:
                d = f"{desc} {nf} images, {nc} corrupt"
                TQDM(None, desc=d, total=n, initial=n)
                if cache["msgs"]:
                    LOGGER.info("\n".join(cache["msgs"]))  # display warnings
            return samples

        # Run scan if *.cache retrieval failed
        nf, nc, msgs, samples, x = 0, 0, [], [], {}
        with ThreadPool(NUM_THREADS) as pool:
            results = pool.imap(func=verify_image, iterable=zip(self.samples, repeat(self.prefix)))
            pbar = TQDM(results, desc=desc, total=len(self.samples))
            for sample, nf_f, nc_f, msg in pbar:
                if nf_f:
                    samples.append(sample)
                if msg:
                    msgs.append(msg)
                nf += nf_f
                nc += nc_f
                pbar.desc = f"{desc} {nf} images, {nc} corrupt"
            pbar.close()
        if msgs:
            LOGGER.info("\n".join(msgs))
        x["hash"] = get_hash([x[0] for x in self.samples])
        x["results"] = nf, nc, len(samples), samples
        x["msgs"] = msgs  # warnings
        save_dataset_cache_file(self.prefix, path, x, DATASET_CACHE_VERSION)
        return samples