Ir para o conteúdo

Como treinar YOLO COCO sem necessidade de conversão

Por que treinar diretamente no COCO

Anotações em COCO o formato pode ser utilizado diretamente para YOLO Ultralytics treinar sem converter para .txt ficheiros primeiro. Isto é feito através da criação de uma subclasse YOLODataset para analisar COCO em tempo real e integrá-lo no fluxo de treino através de um modelo de treino personalizado.

Esta abordagem mantém o COCO como a única fonte de verdade — não convert_coco() chamada, sem reorganização do diretório, sem ficheiros de etiqueta intermédios. YOLO26 e todos os outros modelosYOLO Ultralytics são suportados. Os modelos de segmentação e de pose requerem campos de rótulo adicionais (ver FAQ).

Prefere uma conversão pontual?

Veja o arquivo Guia YOLO COCO YOLO para a norma convert_coco() fluxo de trabalho.

Visão Geral da Arquitetura

São necessárias duas turmas:

  1. COCOJSONDataset — lê ficheiros COCO e converte-os bounding boxes para YOLO na memória durante o treino
  2. COCOJSONTrainer — substituições build_dataset() para usar COCOJSONDataset em vez do padrão YOLODataset

A implementação segue o mesmo padrão que a função integrada GroundingDataset, que também lê diretamente as anotações JSON. São substituídos três métodos: get_img_files(), cache_labels(), e get_labels().

Criação da classe do conjunto de dados COCO

O COCOJSONDataset a classe herda de YOLODataset e substitui a lógica de carregamento das etiquetas. Em vez de ler .txt a partir de ficheiros do diretório «labels», abre o ficheiro COCO , percorre as anotações agrupadas por imagem e converte cada caixa delimitadora do formato COCO [x_min, y_min, width, height] para o formato de centro YOLO [x_center, y_center, width, height]. Anotações colaborativas (iscrowd: 1) e as caixas com área nula são ignoradas automaticamente.

O get_img_files() O método devolve uma lista vazia porque os caminhos das imagens são obtidos a partir do JSON file_name campo interno cache_labels(). Os IDs das categorias são ordenados e remapeados para índices de classe com indexação a zero, pelo que tanto os esquemas de ID com indexação a 1 ( COCO padrão) como os não contíguos funcionam corretamente.

import json
from collections import defaultdict
from pathlib import Path

import numpy as np

from ultralytics.data.dataset import DATASET_CACHE_VERSION, YOLODataset
from ultralytics.data.utils import get_hash, load_dataset_cache_file, save_dataset_cache_file
from ultralytics.utils import TQDM


class COCOJSONDataset(YOLODataset):
    """Dataset that reads COCO JSON annotations directly without conversion to .txt files."""

    def __init__(self, *args, json_file="", **kwargs):
        self.json_file = json_file
        super().__init__(*args, data={"channels": 3}, **kwargs)

    def get_img_files(self, img_path):
        """Image paths are resolved from the JSON file, not from scanning a directory."""
        return []

    def cache_labels(self, path=Path("./labels.cache")):
        """Parse COCO JSON and convert annotations to YOLO format. Results are saved to a .cache file."""
        x = {"labels": []}
        with open(self.json_file) as f:
            coco = json.load(f)

        images = {img["id"]: img for img in coco["images"]}

        # Sort categories by ID and map to 0-indexed classes
        categories = {cat["id"]: i for i, cat in enumerate(sorted(coco["categories"], key=lambda c: c["id"]))}

        img_to_anns = defaultdict(list)
        for ann in coco["annotations"]:
            img_to_anns[ann["image_id"]].append(ann)

        for img_info in TQDM(coco["images"], desc="reading annotations"):
            h, w = img_info["height"], img_info["width"]
            im_file = Path(self.img_path) / img_info["file_name"]
            if not im_file.exists():
                continue

            self.im_files.append(str(im_file))
            bboxes = []
            for ann in img_to_anns.get(img_info["id"], []):
                if ann.get("iscrowd", False):
                    continue
                # COCO: [x, y, w, h] top-left in pixels -> YOLO: [cx, cy, w, h] center normalized
                box = np.array(ann["bbox"], dtype=np.float32)
                box[:2] += box[2:] / 2  # top-left to center
                box[[0, 2]] /= w  # normalize x
                box[[1, 3]] /= h  # normalize y
                if box[2] <= 0 or box[3] <= 0:
                    continue
                cls = categories[ann["category_id"]]
                bboxes.append([cls, *box.tolist()])

            lb = np.array(bboxes, dtype=np.float32) if bboxes else np.zeros((0, 5), dtype=np.float32)
            x["labels"].append(
                {
                    "im_file": str(im_file),
                    "shape": (h, w),
                    "cls": lb[:, 0:1],
                    "bboxes": lb[:, 1:],
                    "segments": [],
                    "normalized": True,
                    "bbox_format": "xywh",
                }
            )
        x["hash"] = get_hash([self.json_file, str(self.img_path)])
        save_dataset_cache_file(self.prefix, path, x, DATASET_CACHE_VERSION)
        return x

    def get_labels(self):
        """Load labels from .cache file if available, otherwise parse JSON and create the cache."""
        cache_path = Path(self.json_file).with_suffix(".cache")
        try:
            cache = load_dataset_cache_file(cache_path)
            assert cache["version"] == DATASET_CACHE_VERSION
            assert cache["hash"] == get_hash([self.json_file, str(self.img_path)])
            self.im_files = [lb["im_file"] for lb in cache["labels"]]
        except (FileNotFoundError, AssertionError, AttributeError, KeyError, ModuleNotFoundError):
            cache = self.cache_labels(cache_path)
        cache.pop("hash", None)
        cache.pop("version", None)
        return cache["labels"]

As etiquetas analisadas são guardadas num .cache ficheiro ao lado do JSON (por exemplo, instances_train.cache). Nas execuções de treino subsequentes, a cache é carregada diretamente, ignorando a análise do JSON. Se o ficheiro JSON for alterado, a verificação do hash falha e a cache é reconstruída automaticamente.

Ligar o conjunto de dados ao pipeline de treino

A única alteração necessária no programa de treino é a substituição build_dataset(). O padrão DetectionTrainer constrói um YOLODataset que procura .txt ficheiros de etiquetas. Substituindo-o por COCOJSONDataset, o treinador lê, em vez disso, o ficheiro COCO do COCO .

O caminho do ficheiro JSON é obtido a partir de um train_json / val_json field in the data config (see Step 3). During training, mode="train" decide train_json; durante a validação, mode="val" decide val_json. Se val_json se não estiver definido, recorre a train_json.

from ultralytics.models.yolo.detect import DetectionTrainer
from ultralytics.utils import colorstr


class COCOJSONTrainer(DetectionTrainer):
    """Trainer that uses COCOJSONDataset for direct COCO JSON training."""

    def build_dataset(self, img_path, mode="train", batch=None):
        json_file = self.data["train_json"] if mode == "train" else self.data.get("val_json", self.data["train_json"])
        return COCOJSONDataset(
            img_path=img_path,
            json_file=json_file,
            imgsz=self.args.imgsz,
            batch_size=batch,
            augment=mode == "train",
            hyp=self.args,
            rect=self.args.rect or mode == "val",
            cache=self.args.cache or None,
            single_cls=self.args.single_cls or False,
            stride=int(self.model.stride.max()) if hasattr(self, "model") and self.model else 32,
            pad=0.0 if mode == "train" else 0.5,
            prefix=colorstr(f"{mode}: "),
            task=self.args.task,
            classes=self.args.classes,
            fraction=self.args.fraction if mode == "train" else 1.0,
        )

Configurar o ficheiro dataset.yaml COCO

O dataset.yaml utiliza o padrão path, train, e val campos para localizar os diretórios de imagens. Dois campos adicionais, train_json e val_json, especifique os ficheiros COCO que COCOJSONTrainer diz. O nc e names Os campos definem o número de classes e os seus nomes, de acordo com a ordem em que aparecem categories no JSON.

path: /path/to/images # root directory with train/ and val/ subfolders
train: train
val: val

# COCO JSON annotation files
train_json: /path/to/annotations/instances_train.json
val_json: /path/to/annotations/instances_val.json

nc: 80
names:
    0: person
    1: bicycle
    # ... remaining class names

Estrutura de diretórios prevista:

my_dataset/
  images/
    train/
      img_001.jpg
      ...
    val/
      img_100.jpg
      ...
  annotations/
    instances_train.json
    instances_val.json
  dataset.yaml

Formação em Running no COCO

Com a classe do conjunto de dados, a classe do treinador e a configuração YAML definidas, o treino decorre de acordo com o padrão model.train() treino. A única diferença em relação a um treino normal é o trainer=COCOJSONTrainer argumento, que indica Ultralytics utilizar o carregador de conjuntos de dados personalizado em vez do carregador predefinido.

from ultralytics import YOLO

model = YOLO("yolo26n.pt")
model.train(data="dataset.yaml", epochs=100, imgsz=640, trainer=COCOJSONTrainer)

Todo o processo de treino decorre conforme o esperado, incluindo a validação, o armazenamento de pontos de verificação e o registo de métricas.

Implementação total

Por uma questão de conveniência, a implementação completa é apresentada abaixo como um único script pronto a copiar e colar. Inclui o conjunto de dados personalizado, o modelo de treino personalizado e a chamada de treino. Guarde isto juntamente com o seu dataset.yaml e executá-lo diretamente.

import json
from collections import defaultdict
from pathlib import Path

import numpy as np

from ultralytics import YOLO
from ultralytics.data.dataset import DATASET_CACHE_VERSION, YOLODataset
from ultralytics.data.utils import get_hash, load_dataset_cache_file, save_dataset_cache_file
from ultralytics.models.yolo.detect import DetectionTrainer
from ultralytics.utils import TQDM, colorstr


class COCOJSONDataset(YOLODataset):
    """Dataset that reads COCO JSON annotations directly without conversion to .txt files."""

    def __init__(self, *args, json_file="", **kwargs):
        self.json_file = json_file
        super().__init__(*args, data={"channels": 3}, **kwargs)

    def get_img_files(self, img_path):
        return []

    def cache_labels(self, path=Path("./labels.cache")):
        x = {"labels": []}
        with open(self.json_file) as f:
            coco = json.load(f)

        images = {img["id"]: img for img in coco["images"]}
        categories = {cat["id"]: i for i, cat in enumerate(sorted(coco["categories"], key=lambda c: c["id"]))}

        img_to_anns = defaultdict(list)
        for ann in coco["annotations"]:
            img_to_anns[ann["image_id"]].append(ann)

        for img_info in TQDM(coco["images"], desc="reading annotations"):
            h, w = img_info["height"], img_info["width"]
            im_file = Path(self.img_path) / img_info["file_name"]
            if not im_file.exists():
                continue

            self.im_files.append(str(im_file))
            bboxes = []
            for ann in img_to_anns.get(img_info["id"], []):
                if ann.get("iscrowd", False):
                    continue
                box = np.array(ann["bbox"], dtype=np.float32)
                box[:2] += box[2:] / 2
                box[[0, 2]] /= w
                box[[1, 3]] /= h
                if box[2] <= 0 or box[3] <= 0:
                    continue
                cls = categories[ann["category_id"]]
                bboxes.append([cls, *box.tolist()])

            lb = np.array(bboxes, dtype=np.float32) if bboxes else np.zeros((0, 5), dtype=np.float32)
            x["labels"].append(
                {
                    "im_file": str(im_file),
                    "shape": (h, w),
                    "cls": lb[:, 0:1],
                    "bboxes": lb[:, 1:],
                    "segments": [],
                    "normalized": True,
                    "bbox_format": "xywh",
                }
            )
        x["hash"] = get_hash([self.json_file, str(self.img_path)])
        save_dataset_cache_file(self.prefix, path, x, DATASET_CACHE_VERSION)
        return x

    def get_labels(self):
        cache_path = Path(self.json_file).with_suffix(".cache")
        try:
            cache = load_dataset_cache_file(cache_path)
            assert cache["version"] == DATASET_CACHE_VERSION
            assert cache["hash"] == get_hash([self.json_file, str(self.img_path)])
            self.im_files = [lb["im_file"] for lb in cache["labels"]]
        except (FileNotFoundError, AssertionError, AttributeError, KeyError, ModuleNotFoundError):
            cache = self.cache_labels(cache_path)
        cache.pop("hash", None)
        cache.pop("version", None)
        return cache["labels"]


class COCOJSONTrainer(DetectionTrainer):
    """Trainer that uses COCOJSONDataset for direct COCO JSON training."""

    def build_dataset(self, img_path, mode="train", batch=None):
        json_file = self.data["train_json"] if mode == "train" else self.data.get("val_json", self.data["train_json"])
        return COCOJSONDataset(
            img_path=img_path,
            json_file=json_file,
            imgsz=self.args.imgsz,
            batch_size=batch,
            augment=mode == "train",
            hyp=self.args,
            rect=self.args.rect or mode == "val",
            cache=self.args.cache or None,
            single_cls=self.args.single_cls or False,
            stride=int(self.model.stride.max()) if hasattr(self, "model") and self.model else 32,
            pad=0.0 if mode == "train" else 0.5,
            prefix=colorstr(f"{mode}: "),
            task=self.args.task,
            classes=self.args.classes,
            fraction=self.args.fraction if mode == "train" else 1.0,
        )


model = YOLO("yolo26n.pt")
model.train(data="dataset.yaml", epochs=100, imgsz=640, trainer=COCOJSONTrainer)

Para obter recomendações sobre hiperparâmetros, consulte o guia «Dicas para o treino de modelos ».

FAQ

Qual é a diferença entre isto e a função convert_coco()?

convert_coco() escreve .txt gravar os ficheiros de anotação no disco como uma conversão única. Esta abordagem analisa o JSON no início de cada execução de treino e converte as anotações na memória. Utilize convert_coco() quando se prefere utilizar rótulos permanentes YOLO; utilize esta abordagem para manter o ficheiro COCO como a única fonte de referência, sem gerar ficheiros adicionais.

YOLO consegue YOLO com COCO sem código personalizado?

Não com o Ultralytics atual Ultralytics , que espera YOLO .txt etiquetas por predefinição. Este guia fornece o código personalizado mínimo necessário — uma classe de conjunto de dados e uma classe de treinador. Uma vez definido, o treino requer apenas um model.train() chamada.

Isto suporta a segmentação e a estimativa de poses?

Este guia aborda detecção de objetos. Para adicionar segmentação de instâncias suporte, inclua o segmentation dados poligonais provenientes das COCO no segments campo de cada dicionário de etiquetas. Para estimativa de pose, incluem keypoints. O GroundingDataset código-fonte fornece uma implementação de referência para o tratamento de segmentos.

As ampliações funcionam com este conjunto de dados personalizado?

Sim. COCOJSONDataset estende-se YOLODataset, pelo que todos os integrados aumentos de dadosmosaic, mixup, copy-paste, entre outros — funcionam sem alterações.

Como é que os IDs das categorias são mapeados para os índices das classes?

As categorias estão ordenadas por id e mapeados para índices sequenciais a partir de 0. Isto permite lidar com IDs com base em 1 ( COCO padrão), IDs com base em 0 e IDs não contíguos. O names dicionário em dataset.yaml deve seguir a mesma ordem de classificação que o COCO categories matriz.

Existe alguma perda de desempenho em comparação com as etiquetas pré-convertidas?

O ficheiro COCO é analisado uma vez durante a primeira execução do treino. Os rótulos analisados são guardados num .cache ficheiro, pelo que as execuções subsequentes carregam instantaneamente sem necessidade de nova análise. A velocidade de treino é idêntica à YOLO padrão YOLO , uma vez que as anotações são mantidas na memória. A cache é reconstruída automaticamente se o ficheiro JSON for alterado.



📅 Criado há 0 dias ✏️ Atualizado há 0 dias
glenn-jocherraimbekovm

Comentários