Saltar al contenido

Referencia para ultralytics/data/dataset.py

Nota

Este archivo est谩 disponible en https://github.com/ultralytics/ ultralytics/blob/main/ ultralytics/data/dataset .py. Si detectas alg煤n problema, por favor, ayuda a solucionarlo contribuyendo con una Pull Request 馃洜锔. 隆Gracias 馃檹!



ultralytics.data.dataset.YOLODataset

Bases: BaseDataset

Clase de conjunto de datos para cargar etiquetas de detecci贸n y/o segmentaci贸n de objetos en formato YOLO .

Par谩metros:

Nombre Tipo Descripci贸n Por defecto
data dict

Un diccionario YAML del conjunto de datos. Por defecto es Ninguno.

None
task str

Un arg expl铆cito para se帽alar la tarea actual, Por defecto es 'detectar'.

'detect'

Devuelve:

Tipo Descripci贸n
Dataset

Un objeto del conjunto de datos PyTorch que puede utilizarse para entrenar un modelo de detecci贸n de objetos.

C贸digo fuente en ultralytics/data/dataset.py
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=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", "segments", "obb"}:
                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

__init__(*args, data=None, task='detect', **kwargs)

Inicializa el conjunto de datos YOLOD con configuraciones opcionales para segmentos y puntos clave.

C贸digo fuente en ultralytics/data/dataset.py
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)

build_transforms(hyp=None)

Construye y a帽ade transformaciones a la lista.

C贸digo fuente en ultralytics/data/dataset.py
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=hyp.bgr if self.augment else 0.0,  # only affect training.
        )
    )
    return transforms

cache_labels(path=Path('./labels.cache'))

Almacena en cach茅 las etiquetas de los conjuntos de datos, comprueba las im谩genes y lee las formas.

Par谩metros:

Nombre Tipo Descripci贸n Por defecto
path Path

Ruta donde guardar el archivo de cach茅. Por defecto es Path('./etiquetas.cache').

Path('./labels.cache')

Devuelve:

Tipo Descripci贸n
dict

etiquetas.

C贸digo fuente en ultralytics/data/dataset.py
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

close_mosaic(hyp)

Establece las opciones mosaico, copiar_pegar y mezcla a 0,0 y construye las transformaciones.

C贸digo fuente en ultralytics/data/dataset.py
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)

collate_fn(batch) staticmethod

Coteja las muestras de datos en lotes.

C贸digo fuente en ultralytics/data/dataset.py
@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", "segments", "obb"}:
            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

get_labels()

Devuelve el diccionario de etiquetas para el entrenamiento YOLO .

C贸digo fuente en ultralytics/data/dataset.py
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

update_labels_info(label)

Personaliza aqu铆 el formato de tus etiquetas.

Nota

cls no est谩 con bboxes ahora, la clasificaci贸n y la segmentaci贸n sem谩ntica necesitan una etiqueta cls independiente Tambi茅n puede soportar la clasificaci贸n y la segmentaci贸n sem谩ntica a帽adiendo o eliminando claves dict all铆.

C贸digo fuente en ultralytics/data/dataset.py
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



ultralytics.data.dataset.YOLOMultiModalDataset

Bases: YOLODataset

Clase de conjunto de datos para cargar etiquetas de detecci贸n y/o segmentaci贸n de objetos en formato YOLO .

Par谩metros:

Nombre Tipo Descripci贸n Por defecto
data dict

Un diccionario YAML del conjunto de datos. Por defecto es Ninguno.

None
task str

Un arg expl铆cito para se帽alar la tarea actual, Por defecto es 'detectar'.

'detect'

Devuelve:

Tipo Descripci贸n
Dataset

Un objeto del conjunto de datos PyTorch que puede utilizarse para entrenar un modelo de detecci贸n de objetos.

C贸digo fuente en ultralytics/data/dataset.py
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

__init__(*args, data=None, task='detect', **kwargs)

Inicializa un objeto de conjunto de datos para tareas de detecci贸n de objetos con especificaciones opcionales.

C贸digo fuente en ultralytics/data/dataset.py
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)

build_transforms(hyp=None)

Mejora las transformaciones de datos con el aumento opcional de texto para la formaci贸n multimodal.

C贸digo fuente en ultralytics/data/dataset.py
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

update_labels_info(label)

A帽ade informaci贸n de textos para el entrenamiento de modelos multimodales.

C贸digo fuente en ultralytics/data/dataset.py
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



ultralytics.data.dataset.GroundingDataset

Bases: YOLODataset

C贸digo fuente en ultralytics/data/dataset.py
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

__init__(*args, task='detect', json_file, **kwargs)

Inicializa un GroundingDataset para la detecci贸n de objetos, cargando anotaciones de un archivo JSON especificado.

C贸digo fuente en ultralytics/data/dataset.py
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)

build_transforms(hyp=None)

Configura los aumentos para el entrenamiento con carga de texto opcional; hyp ajusta la intensidad del aumento.

C贸digo fuente en ultralytics/data/dataset.py
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

get_img_files(img_path)

Los archivos de imagen se leer铆an en get_labels devuelve aqu铆 una lista vac铆a.

C贸digo fuente en ultralytics/data/dataset.py
def get_img_files(self, img_path):
    """The image files would be read in `get_labels` function, return empty list here."""
    return []

get_labels()

Carga las anotaciones desde un archivo JSON, las filtra y normaliza los cuadros delimitadores de cada imagen.

C贸digo fuente en ultralytics/data/dataset.py
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



ultralytics.data.dataset.YOLOConcatDataset

Bases: ConcatDataset

Conjunto de datos como concatenaci贸n de varios conjuntos de datos.

Esta clase es 煤til para reunir diferentes conjuntos de datos existentes.

C贸digo fuente en ultralytics/data/dataset.py
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)

collate_fn(batch) staticmethod

Coteja las muestras de datos en lotes.

C贸digo fuente en ultralytics/data/dataset.py
@staticmethod
def collate_fn(batch):
    """Collates data samples into batches."""
    return YOLODataset.collate_fn(batch)



ultralytics.data.dataset.SemanticDataset

Bases: BaseDataset

Conjunto de datos de segmentaci贸n sem谩ntica.

Esta clase se encarga de manejar los conjuntos de datos utilizados en las tareas de segmentaci贸n sem谩ntica. Hereda funcionalidades de la clase BaseDataset.

Nota

Esta clase es actualmente un marcador de posici贸n y necesita ser rellenada con m茅todos y atributos para apoyar tareas de segmentaci贸n sem谩ntica.

C贸digo fuente en ultralytics/data/dataset.py
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__()

__init__()

Inicializa un objeto SemanticDataset.

C贸digo fuente en ultralytics/data/dataset.py
def __init__(self):
    """Initialize a SemanticDataset object."""
    super().__init__()



ultralytics.data.dataset.ClassificationDataset

Ampl铆a ImageFolder de torchvision para que admita las tareas de clasificaci贸n de YOLO , ofreciendo funcionalidades como imagen, almacenamiento en cach茅 y verificaci贸n. Est谩 dise帽ado para manejar eficientemente grandes conjuntos de datos para el entrenamiento de modelos de aprendizaje profundo, con transformaciones de imagen opcionales y mecanismos de almacenamiento en cach茅 para acelerar el entrenamiento. con transformaciones de imagen opcionales y mecanismos de almacenamiento en cach茅 para acelerar el entrenamiento.

Esta clase permite aumentos utilizando las bibliotecas torchvision y Albumentations, y admite el almacenamiento en cach茅 de im谩genes en RAM o en disco para reducir la sobrecarga de E/S durante el entrenamiento. Adem谩s, implementa un s贸lido proceso de verificaci贸n para garantizar la integridad y coherencia de los datos.

Atributos:

Nombre Tipo Descripci贸n
cache_ram bool

Indica si est谩 activada la cach茅 en RAM.

cache_disk bool

Indica si est谩 activada la cach茅 en disco.

samples list

Una lista de tuplas, cada una de las cuales contiene la ruta a una imagen, su 铆ndice de clase, la ruta a su archivo de cach茅 .npy (si se almacena en disco) y, opcionalmente, la matriz de im谩genes cargada (si se almacena en RAM).

torch_transforms callable

PyTorch transformaciones que se aplicar谩n a las im谩genes.

C贸digo fuente en ultralytics/data/dataset.py
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

__getitem__(i)

Devuelve el subconjunto de datos y objetivos correspondientes a los 铆ndices dados.

C贸digo fuente en ultralytics/data/dataset.py
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}

__init__(root, args, augment=False, prefix='')

Inicializa el objeto YOLO con los ajustes de ra铆z, tama帽o de imagen, aumentos y cach茅.

Par谩metros:

Nombre Tipo Descripci贸n Por defecto
root str

Ruta al directorio del conjunto de datos donde se almacenan las im谩genes en una estructura de carpetas espec铆fica de la clase.

necesario
args Namespace

Configuraci贸n que contiene ajustes relacionados con el conjunto de datos, como el tama帽o de la imagen, el aumento y la configuraci贸n de la cach茅. Incluye atributos como imgsz (tama帽o de la imagen), fraction (fracci贸n de datos a utilizar), scale, fliplr, flipud, cache (cach茅 en disco o RAM para un entrenamiento m谩s r谩pido), auto_augment, hsv_h, hsv_s, hsv_vy crop_fraction.

necesario
augment bool

Si se aplican aumentos al conjunto de datos. Por defecto es Falso.

False
prefix str

Prefijo para los nombres de archivos de registro y cach茅, que ayuda a identificar y depurar los conjuntos de datos. depuraci贸n. Por defecto es una cadena vac铆a.

''
C贸digo fuente en ultralytics/data/dataset.py
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)
    )

__len__()

Devuelve el n煤mero total de muestras del conjunto de datos.

C贸digo fuente en ultralytics/data/dataset.py
def __len__(self) -> int:
    """Return the total number of samples in the dataset."""
    return len(self.samples)

verify_images()

Verifica todas las im谩genes del conjunto de datos.

C贸digo fuente en ultralytics/data/dataset.py
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





Creado 2023-11-12, Actualizado 2024-05-18
Autores: glenn-jocher (5), Burhan-Q (1), Laughing-q (1)