跳至内容

参考资料 ultralytics/data/base.py

备注

该文件可从https://github.com/ultralytics/ultralytics/blob/main/ ultralytics/data/base .py 获取。如果您发现问题,请通过提交 Pull Request🛠️ 帮助修复。谢谢🙏!



ultralytics.data.base.BaseDataset

垒球 Dataset

用于加载和处理图像数据的基础数据集类。

参数

名称 类型 说明 默认值
img_path str

包含图像的文件夹的路径。

所需
imgsz int

图像大小。默认为 640。

640
cache bool

在训练期间将图像缓存到 RAM 或磁盘。默认为 "假"。

False
augment bool

如果为 True,则应用数据增强。默认为 True。

True
hyp dict

用于数据增强的超参数。默认为 "无"。

DEFAULT_CFG
prefix str

在日志信息中打印的前缀。默认为''。

''
rect bool

如果为 True,则使用矩形训练。默认为 "假"。

False
batch_size int

批次大小。默认为 "无"。

16
stride int

步长。默认为 32。

32
pad float

填充。默认为 0.0。

0.5
single_cls bool

如果为 True,则使用单班训练。默认为 "假"。

False
classes list

包含的类的列表。默认为无。

None
fraction float

要使用的数据集的比例。默认为 1.0(使用所有数据)。

1.0

属性

名称 类型 说明
im_files list

图像文件路径列表。

labels list

标签数据字典列表。

ni int

数据集中的图像数量。

ims list

已加载图像的列表。

npy_files list

numpy 文件路径列表。

transforms callable

图像转换功能。

源代码 ultralytics/data/base.py
class BaseDataset(Dataset):
    """
    Base dataset class for loading and processing image data.

    Args:
        img_path (str): Path to the folder containing images.
        imgsz (int, optional): Image size. Defaults to 640.
        cache (bool, optional): Cache images to RAM or disk during training. Defaults to False.
        augment (bool, optional): If True, data augmentation is applied. Defaults to True.
        hyp (dict, optional): Hyperparameters to apply data augmentation. Defaults to None.
        prefix (str, optional): Prefix to print in log messages. Defaults to ''.
        rect (bool, optional): If True, rectangular training is used. Defaults to False.
        batch_size (int, optional): Size of batches. Defaults to None.
        stride (int, optional): Stride. Defaults to 32.
        pad (float, optional): Padding. Defaults to 0.0.
        single_cls (bool, optional): If True, single class training is used. Defaults to False.
        classes (list): List of included classes. Default is None.
        fraction (float): Fraction of dataset to utilize. Default is 1.0 (use all data).

    Attributes:
        im_files (list): List of image file paths.
        labels (list): List of label data dictionaries.
        ni (int): Number of images in the dataset.
        ims (list): List of loaded images.
        npy_files (list): List of numpy file paths.
        transforms (callable): Image transformation function.
    """

    def __init__(
        self,
        img_path,
        imgsz=640,
        cache=False,
        augment=True,
        hyp=DEFAULT_CFG,
        prefix="",
        rect=False,
        batch_size=16,
        stride=32,
        pad=0.5,
        single_cls=False,
        classes=None,
        fraction=1.0,
    ):
        """Initialize BaseDataset with given configuration and options."""
        super().__init__()
        self.img_path = img_path
        self.imgsz = imgsz
        self.augment = augment
        self.single_cls = single_cls
        self.prefix = prefix
        self.fraction = fraction
        self.im_files = self.get_img_files(self.img_path)
        self.labels = self.get_labels()
        self.update_labels(include_class=classes)  # single_cls and include_class
        self.ni = len(self.labels)  # number of images
        self.rect = rect
        self.batch_size = batch_size
        self.stride = stride
        self.pad = pad
        if self.rect:
            assert self.batch_size is not None
            self.set_rectangle()

        # Buffer thread for mosaic images
        self.buffer = []  # buffer size = batch size
        self.max_buffer_length = min((self.ni, self.batch_size * 8, 1000)) if self.augment else 0

        # Cache images
        if cache == "ram" and not self.check_cache_ram():
            cache = False
        self.ims, self.im_hw0, self.im_hw = [None] * self.ni, [None] * self.ni, [None] * self.ni
        self.npy_files = [Path(f).with_suffix(".npy") for f in self.im_files]
        if cache:
            self.cache_images(cache)

        # Transforms
        self.transforms = self.build_transforms(hyp=hyp)

    def get_img_files(self, img_path):
        """Read image files."""
        try:
            f = []  # image files
            for p in img_path if isinstance(img_path, list) else [img_path]:
                p = Path(p)  # os-agnostic
                if p.is_dir():  # dir
                    f += glob.glob(str(p / "**" / "*.*"), recursive=True)
                    # F = list(p.rglob('*.*'))  # pathlib
                elif p.is_file():  # file
                    with open(p) as t:
                        t = t.read().strip().splitlines()
                        parent = str(p.parent) + os.sep
                        f += [x.replace("./", parent) if x.startswith("./") else x for x in t]  # local to global path
                        # F += [p.parent / x.lstrip(os.sep) for x in t]  # local to global path (pathlib)
                else:
                    raise FileNotFoundError(f"{self.prefix}{p} does not exist")
            im_files = sorted(x.replace("/", os.sep) for x in f if x.split(".")[-1].lower() in IMG_FORMATS)
            # self.img_files = sorted([x for x in f if x.suffix[1:].lower() in IMG_FORMATS])  # pathlib
            assert im_files, f"{self.prefix}No images found in {img_path}"
        except Exception as e:
            raise FileNotFoundError(f"{self.prefix}Error loading data from {img_path}\n{HELP_URL}") from e
        if self.fraction < 1:
            im_files = im_files[: round(len(im_files) * self.fraction)]
        return im_files

    def update_labels(self, include_class: Optional[list]):
        """Update labels to include only these classes (optional)."""
        include_class_array = np.array(include_class).reshape(1, -1)
        for i in range(len(self.labels)):
            if include_class is not None:
                cls = self.labels[i]["cls"]
                bboxes = self.labels[i]["bboxes"]
                segments = self.labels[i]["segments"]
                keypoints = self.labels[i]["keypoints"]
                j = (cls == include_class_array).any(1)
                self.labels[i]["cls"] = cls[j]
                self.labels[i]["bboxes"] = bboxes[j]
                if segments:
                    self.labels[i]["segments"] = [segments[si] for si, idx in enumerate(j) if idx]
                if keypoints is not None:
                    self.labels[i]["keypoints"] = keypoints[j]
            if self.single_cls:
                self.labels[i]["cls"][:, 0] = 0

    def load_image(self, i, rect_mode=True):
        """Loads 1 image from dataset index 'i', returns (im, resized hw)."""
        im, f, fn = self.ims[i], self.im_files[i], self.npy_files[i]
        if im is None:  # not cached in RAM
            if fn.exists():  # load npy
                try:
                    im = np.load(fn)
                except Exception as e:
                    LOGGER.warning(f"{self.prefix}WARNING ⚠️ Removing corrupt *.npy image file {fn} due to: {e}")
                    Path(fn).unlink(missing_ok=True)
                    im = cv2.imread(f)  # BGR
            else:  # read image
                im = cv2.imread(f)  # BGR
            if im is None:
                raise FileNotFoundError(f"Image Not Found {f}")

            h0, w0 = im.shape[:2]  # orig hw
            if rect_mode:  # resize long side to imgsz while maintaining aspect ratio
                r = self.imgsz / max(h0, w0)  # ratio
                if r != 1:  # if sizes are not equal
                    w, h = (min(math.ceil(w0 * r), self.imgsz), min(math.ceil(h0 * r), self.imgsz))
                    im = cv2.resize(im, (w, h), interpolation=cv2.INTER_LINEAR)
            elif not (h0 == w0 == self.imgsz):  # resize by stretching image to square imgsz
                im = cv2.resize(im, (self.imgsz, self.imgsz), interpolation=cv2.INTER_LINEAR)

            # Add to buffer if training with augmentations
            if self.augment:
                self.ims[i], self.im_hw0[i], self.im_hw[i] = im, (h0, w0), im.shape[:2]  # im, hw_original, hw_resized
                self.buffer.append(i)
                if len(self.buffer) >= self.max_buffer_length:
                    j = self.buffer.pop(0)
                    self.ims[j], self.im_hw0[j], self.im_hw[j] = None, None, None

            return im, (h0, w0), im.shape[:2]

        return self.ims[i], self.im_hw0[i], self.im_hw[i]

    def cache_images(self, cache):
        """Cache images to memory or disk."""
        b, gb = 0, 1 << 30  # bytes of cached images, bytes per gigabytes
        fcn = self.cache_images_to_disk if cache == "disk" else self.load_image
        with ThreadPool(NUM_THREADS) as pool:
            results = pool.imap(fcn, range(self.ni))
            pbar = TQDM(enumerate(results), total=self.ni, disable=LOCAL_RANK > 0)
            for i, x in pbar:
                if cache == "disk":
                    b += self.npy_files[i].stat().st_size
                else:  # 'ram'
                    self.ims[i], self.im_hw0[i], self.im_hw[i] = x  # im, hw_orig, hw_resized = load_image(self, i)
                    b += self.ims[i].nbytes
                pbar.desc = f"{self.prefix}Caching images ({b / gb:.1f}GB {cache})"
            pbar.close()

    def cache_images_to_disk(self, i):
        """Saves an image as an *.npy file for faster loading."""
        f = self.npy_files[i]
        if not f.exists():
            np.save(f.as_posix(), cv2.imread(self.im_files[i]), allow_pickle=False)

    def check_cache_ram(self, safety_margin=0.5):
        """Check image caching requirements vs available memory."""
        b, gb = 0, 1 << 30  # bytes of cached images, bytes per gigabytes
        n = min(self.ni, 30)  # extrapolate from 30 random images
        for _ in range(n):
            im = cv2.imread(random.choice(self.im_files))  # sample image
            ratio = self.imgsz / max(im.shape[0], im.shape[1])  # max(h, w)  # ratio
            b += im.nbytes * ratio**2
        mem_required = b * self.ni / n * (1 + safety_margin)  # GB required to cache dataset into RAM
        mem = psutil.virtual_memory()
        cache = mem_required < mem.available  # to cache or not to cache, that is the question
        if not cache:
            LOGGER.info(
                f'{self.prefix}{mem_required / gb:.1f}GB RAM required to cache images '
                f'with {int(safety_margin * 100)}% safety margin but only '
                f'{mem.available / gb:.1f}/{mem.total / gb:.1f}GB available, '
                f"{'caching images ✅' if cache else 'not caching images ⚠️'}"
            )
        return cache

    def set_rectangle(self):
        """Sets the shape of bounding boxes for YOLO detections as rectangles."""
        bi = np.floor(np.arange(self.ni) / self.batch_size).astype(int)  # batch index
        nb = bi[-1] + 1  # number of batches

        s = np.array([x.pop("shape") for x in self.labels])  # hw
        ar = s[:, 0] / s[:, 1]  # aspect ratio
        irect = ar.argsort()
        self.im_files = [self.im_files[i] for i in irect]
        self.labels = [self.labels[i] for i in irect]
        ar = ar[irect]

        # Set training image shapes
        shapes = [[1, 1]] * nb
        for i in range(nb):
            ari = ar[bi == i]
            mini, maxi = ari.min(), ari.max()
            if maxi < 1:
                shapes[i] = [maxi, 1]
            elif mini > 1:
                shapes[i] = [1, 1 / mini]

        self.batch_shapes = np.ceil(np.array(shapes) * self.imgsz / self.stride + self.pad).astype(int) * self.stride
        self.batch = bi  # batch index of image

    def __getitem__(self, index):
        """Returns transformed label information for given index."""
        return self.transforms(self.get_image_and_label(index))

    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]]
        return self.update_labels_info(label)

    def __len__(self):
        """Returns the length of the labels list for the dataset."""
        return len(self.labels)

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

    def build_transforms(self, hyp=None):
        """
        Users can customize augmentations here.

        Example:
            ```python
            if self.augment:
                # Training transforms
                return Compose([])
            else:
                # Val transforms
                return Compose([])
            ```
        """
        raise NotImplementedError

    def get_labels(self):
        """
        Users can customize their own format here.

        Note:
            Ensure output is a dictionary with the following keys:
            ```python
            dict(
                im_file=im_file,
                shape=shape,  # format: (height, width)
                cls=cls,
                bboxes=bboxes, # xywh
                segments=segments,  # xy
                keypoints=keypoints, # xy
                normalized=True, # or False
                bbox_format="xyxy",  # or xywh, ltwh
            )
            ```
        """
        raise NotImplementedError

__getitem__(index)

返回给定索引的转换标签信息。

源代码 ultralytics/data/base.py
def __getitem__(self, index):
    """Returns transformed label information for given index."""
    return self.transforms(self.get_image_and_label(index))

__init__(img_path, imgsz=640, cache=False, augment=True, hyp=DEFAULT_CFG, prefix='', rect=False, batch_size=16, stride=32, pad=0.5, single_cls=False, classes=None, fraction=1.0)

使用给定的配置和选项初始化 BaseDataset。

源代码 ultralytics/data/base.py
def __init__(
    self,
    img_path,
    imgsz=640,
    cache=False,
    augment=True,
    hyp=DEFAULT_CFG,
    prefix="",
    rect=False,
    batch_size=16,
    stride=32,
    pad=0.5,
    single_cls=False,
    classes=None,
    fraction=1.0,
):
    """Initialize BaseDataset with given configuration and options."""
    super().__init__()
    self.img_path = img_path
    self.imgsz = imgsz
    self.augment = augment
    self.single_cls = single_cls
    self.prefix = prefix
    self.fraction = fraction
    self.im_files = self.get_img_files(self.img_path)
    self.labels = self.get_labels()
    self.update_labels(include_class=classes)  # single_cls and include_class
    self.ni = len(self.labels)  # number of images
    self.rect = rect
    self.batch_size = batch_size
    self.stride = stride
    self.pad = pad
    if self.rect:
        assert self.batch_size is not None
        self.set_rectangle()

    # Buffer thread for mosaic images
    self.buffer = []  # buffer size = batch size
    self.max_buffer_length = min((self.ni, self.batch_size * 8, 1000)) if self.augment else 0

    # Cache images
    if cache == "ram" and not self.check_cache_ram():
        cache = False
    self.ims, self.im_hw0, self.im_hw = [None] * self.ni, [None] * self.ni, [None] * self.ni
    self.npy_files = [Path(f).with_suffix(".npy") for f in self.im_files]
    if cache:
        self.cache_images(cache)

    # Transforms
    self.transforms = self.build_transforms(hyp=hyp)

__len__()

返回数据集标签列表的长度。

源代码 ultralytics/data/base.py
def __len__(self):
    """Returns the length of the labels list for the dataset."""
    return len(self.labels)

build_transforms(hyp=None)

用户可以在这里定制增强功能。

示例
if self.augment:
    # Training transforms
    return Compose([])
else:
    # Val transforms
    return Compose([])
源代码 ultralytics/data/base.py
def build_transforms(self, hyp=None):
    """
    Users can customize augmentations here.

    Example:
        ```python
        if self.augment:
            # Training transforms
            return Compose([])
        else:
            # Val transforms
            return Compose([])
        ```
    """
    raise NotImplementedError

cache_images(cache)

将图像缓存到内存或磁盘。

源代码 ultralytics/data/base.py
def cache_images(self, cache):
    """Cache images to memory or disk."""
    b, gb = 0, 1 << 30  # bytes of cached images, bytes per gigabytes
    fcn = self.cache_images_to_disk if cache == "disk" else self.load_image
    with ThreadPool(NUM_THREADS) as pool:
        results = pool.imap(fcn, range(self.ni))
        pbar = TQDM(enumerate(results), total=self.ni, disable=LOCAL_RANK > 0)
        for i, x in pbar:
            if cache == "disk":
                b += self.npy_files[i].stat().st_size
            else:  # 'ram'
                self.ims[i], self.im_hw0[i], self.im_hw[i] = x  # im, hw_orig, hw_resized = load_image(self, i)
                b += self.ims[i].nbytes
            pbar.desc = f"{self.prefix}Caching images ({b / gb:.1f}GB {cache})"
        pbar.close()

cache_images_to_disk(i)

将图像保存为 *.npy 文件,以加快加载速度。

源代码 ultralytics/data/base.py
def cache_images_to_disk(self, i):
    """Saves an image as an *.npy file for faster loading."""
    f = self.npy_files[i]
    if not f.exists():
        np.save(f.as_posix(), cv2.imread(self.im_files[i]), allow_pickle=False)

check_cache_ram(safety_margin=0.5)

检查图像缓存要求和可用内存。

源代码 ultralytics/data/base.py
def check_cache_ram(self, safety_margin=0.5):
    """Check image caching requirements vs available memory."""
    b, gb = 0, 1 << 30  # bytes of cached images, bytes per gigabytes
    n = min(self.ni, 30)  # extrapolate from 30 random images
    for _ in range(n):
        im = cv2.imread(random.choice(self.im_files))  # sample image
        ratio = self.imgsz / max(im.shape[0], im.shape[1])  # max(h, w)  # ratio
        b += im.nbytes * ratio**2
    mem_required = b * self.ni / n * (1 + safety_margin)  # GB required to cache dataset into RAM
    mem = psutil.virtual_memory()
    cache = mem_required < mem.available  # to cache or not to cache, that is the question
    if not cache:
        LOGGER.info(
            f'{self.prefix}{mem_required / gb:.1f}GB RAM required to cache images '
            f'with {int(safety_margin * 100)}% safety margin but only '
            f'{mem.available / gb:.1f}/{mem.total / gb:.1f}GB available, '
            f"{'caching images ✅' if cache else 'not caching images ⚠️'}"
        )
    return cache

get_image_and_label(index)

从数据集中获取并返回标签信息。

源代码 ultralytics/data/base.py
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]]
    return self.update_labels_info(label)

get_img_files(img_path)

读取图像文件

源代码 ultralytics/data/base.py
def get_img_files(self, img_path):
    """Read image files."""
    try:
        f = []  # image files
        for p in img_path if isinstance(img_path, list) else [img_path]:
            p = Path(p)  # os-agnostic
            if p.is_dir():  # dir
                f += glob.glob(str(p / "**" / "*.*"), recursive=True)
                # F = list(p.rglob('*.*'))  # pathlib
            elif p.is_file():  # file
                with open(p) as t:
                    t = t.read().strip().splitlines()
                    parent = str(p.parent) + os.sep
                    f += [x.replace("./", parent) if x.startswith("./") else x for x in t]  # local to global path
                    # F += [p.parent / x.lstrip(os.sep) for x in t]  # local to global path (pathlib)
            else:
                raise FileNotFoundError(f"{self.prefix}{p} does not exist")
        im_files = sorted(x.replace("/", os.sep) for x in f if x.split(".")[-1].lower() in IMG_FORMATS)
        # self.img_files = sorted([x for x in f if x.suffix[1:].lower() in IMG_FORMATS])  # pathlib
        assert im_files, f"{self.prefix}No images found in {img_path}"
    except Exception as e:
        raise FileNotFoundError(f"{self.prefix}Error loading data from {img_path}\n{HELP_URL}") from e
    if self.fraction < 1:
        im_files = im_files[: round(len(im_files) * self.fraction)]
    return im_files

get_labels()

用户可以在这里定制自己的格式。

备注

确保输出是一个包含以下键值的字典:

dict(
    im_file=im_file,
    shape=shape,  # format: (height, width)
    cls=cls,
    bboxes=bboxes, # xywh
    segments=segments,  # xy
    keypoints=keypoints, # xy
    normalized=True, # or False
    bbox_format="xyxy",  # or xywh, ltwh
)

源代码 ultralytics/data/base.py
def get_labels(self):
    """
    Users can customize their own format here.

    Note:
        Ensure output is a dictionary with the following keys:
        ```python
        dict(
            im_file=im_file,
            shape=shape,  # format: (height, width)
            cls=cls,
            bboxes=bboxes, # xywh
            segments=segments,  # xy
            keypoints=keypoints, # xy
            normalized=True, # or False
            bbox_format="xyxy",  # or xywh, ltwh
        )
        ```
    """
    raise NotImplementedError

load_image(i, rect_mode=True)

从索引 "i "的数据集中加载 1 幅图像,返回(im,调整后的 hw)。

源代码 ultralytics/data/base.py
def load_image(self, i, rect_mode=True):
    """Loads 1 image from dataset index 'i', returns (im, resized hw)."""
    im, f, fn = self.ims[i], self.im_files[i], self.npy_files[i]
    if im is None:  # not cached in RAM
        if fn.exists():  # load npy
            try:
                im = np.load(fn)
            except Exception as e:
                LOGGER.warning(f"{self.prefix}WARNING ⚠️ Removing corrupt *.npy image file {fn} due to: {e}")
                Path(fn).unlink(missing_ok=True)
                im = cv2.imread(f)  # BGR
        else:  # read image
            im = cv2.imread(f)  # BGR
        if im is None:
            raise FileNotFoundError(f"Image Not Found {f}")

        h0, w0 = im.shape[:2]  # orig hw
        if rect_mode:  # resize long side to imgsz while maintaining aspect ratio
            r = self.imgsz / max(h0, w0)  # ratio
            if r != 1:  # if sizes are not equal
                w, h = (min(math.ceil(w0 * r), self.imgsz), min(math.ceil(h0 * r), self.imgsz))
                im = cv2.resize(im, (w, h), interpolation=cv2.INTER_LINEAR)
        elif not (h0 == w0 == self.imgsz):  # resize by stretching image to square imgsz
            im = cv2.resize(im, (self.imgsz, self.imgsz), interpolation=cv2.INTER_LINEAR)

        # Add to buffer if training with augmentations
        if self.augment:
            self.ims[i], self.im_hw0[i], self.im_hw[i] = im, (h0, w0), im.shape[:2]  # im, hw_original, hw_resized
            self.buffer.append(i)
            if len(self.buffer) >= self.max_buffer_length:
                j = self.buffer.pop(0)
                self.ims[j], self.im_hw0[j], self.im_hw[j] = None, None, None

        return im, (h0, w0), im.shape[:2]

    return self.ims[i], self.im_hw0[i], self.im_hw[i]

set_rectangle()

将YOLO 检测的边界框设置为矩形。

源代码 ultralytics/data/base.py
def set_rectangle(self):
    """Sets the shape of bounding boxes for YOLO detections as rectangles."""
    bi = np.floor(np.arange(self.ni) / self.batch_size).astype(int)  # batch index
    nb = bi[-1] + 1  # number of batches

    s = np.array([x.pop("shape") for x in self.labels])  # hw
    ar = s[:, 0] / s[:, 1]  # aspect ratio
    irect = ar.argsort()
    self.im_files = [self.im_files[i] for i in irect]
    self.labels = [self.labels[i] for i in irect]
    ar = ar[irect]

    # Set training image shapes
    shapes = [[1, 1]] * nb
    for i in range(nb):
        ari = ar[bi == i]
        mini, maxi = ari.min(), ari.max()
        if maxi < 1:
            shapes[i] = [maxi, 1]
        elif mini > 1:
            shapes[i] = [1, 1 / mini]

    self.batch_shapes = np.ceil(np.array(shapes) * self.imgsz / self.stride + self.pad).astype(int) * self.stride
    self.batch = bi  # batch index of image

update_labels(include_class)

更新标签,只包含这些类别(可选)。

源代码 ultralytics/data/base.py
def update_labels(self, include_class: Optional[list]):
    """Update labels to include only these classes (optional)."""
    include_class_array = np.array(include_class).reshape(1, -1)
    for i in range(len(self.labels)):
        if include_class is not None:
            cls = self.labels[i]["cls"]
            bboxes = self.labels[i]["bboxes"]
            segments = self.labels[i]["segments"]
            keypoints = self.labels[i]["keypoints"]
            j = (cls == include_class_array).any(1)
            self.labels[i]["cls"] = cls[j]
            self.labels[i]["bboxes"] = bboxes[j]
            if segments:
                self.labels[i]["segments"] = [segments[si] for si, idx in enumerate(j) if idx]
            if keypoints is not None:
                self.labels[i]["keypoints"] = keypoints[j]
        if self.single_cls:
            self.labels[i]["cls"][:, 0] = 0

update_labels_info(label)

在此自定义标签格式。

源代码 ultralytics/data/base.py
def update_labels_info(self, label):
    """Custom your label format here."""
    return label





创建于 2023-11-12,更新于 2023-11-25
作者:glenn-jocher(3)