1 基础环境安装

见torch镜像制作昇腾-Ubuntu镜像制作_乌班图25.1百度云-CSDN博客。

2 安装ultralytics

pip install ultralytics opencv-python-headless onnx -i https://pypi.tuna.tsinghua.edu.cn/simple
apt-get install libgl1
apt-get install libglib2.0-0

3 代码迁移

torch_utils.py

vim /root/miniconda3/envs/python311/lib/python3.11/site-packages/ultralytics/utils/torch_utils.py

# 代码开始部分添加torch_npu
import torch
import torch_npu

# 大约220行开始，修改有关cuda检查的代码

        if "," in device:
            device = ",".join([x for x in device.split(",") if x])  # remove sequential commas, i.e. "0,,1" -> "0,1"
        visible = os.environ.get("CUDA_VISIBLE_DEVICES", None)
        os.environ["CUDA_VISIBLE_DEVICES"] = device  # set environment variable - must be before assert is_available()
        # if not (torch.cuda.is_available() and torch.cuda.device_count() >= len(device.split(","))):
        if not (torch.npu.is_available() and torch.npu.device_count() >= len(device.split(","))):
            LOGGER.info(s)
            install = (
                "See https://pytorch.org/get-started/locally/ for up-to-date torch install instructions if no "
                "CUDA devices are seen by torch.\n"
                # if torch.cuda.device_count() == 0
                if torch.npu.device_count() == 0
                else ""
            )
            raise ValueError(
                f"Invalid CUDA 'device={device}' requested."
                f" Use 'device=cpu' or pass valid CUDA device(s) if available,"
                f" i.e. 'device=0' or 'device=0,1,2,3' for Multi-GPU.\n"
                f"\ntorch.cuda.is_available(): {torch.cuda.is_available()}"
                f"\ntorch.cuda.device_count(): {torch.cuda.device_count()}"
                f"\nos.environ['CUDA_VISIBLE_DEVICES']: {visible}\n"
                f"{install}"
            )

4 数据集准备

4.1 下载数据集

COCO2017数据集下载并解压（可以先解压annotations_trainval2017.zip，图片留到后面构造数据集目录时再解压到指定的目录下）

modelscope download --dataset PAI/COCO2017 --local_dir ./COCO2017

解压后的目录**（多的json文件要删掉，否则4.2节会有报错）**

├── COCO2017
│   ├── annotations
│   ├── annotations_trainval2017.zip
│   ├── COCO2017.json
│   ├── README.md
│   ├── train2017
│   ├── train2017.zip
│   ├── val2017
│   └── val2017.zip

4.2 数据集格式转换

utils.py、coco_dataset.py在同一级目录下

4.2.1 utils.py

# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license

import glob
import os
import shutil
from pathlib import Path

import numpy as np
from PIL import ExifTags
from tqdm import tqdm

# Parameters
img_formats = ["bmp", "jpg", "jpeg", "png", "tif", "tiff", "dng"]  # acceptable image suffixes
vid_formats = ["mov", "avi", "mp4", "mpg", "mpeg", "m4v", "wmv", "mkv"]  # acceptable video suffixes

# Get orientation exif tag
for orientation in ExifTags.TAGS.keys():
    if ExifTags.TAGS[orientation] == "Orientation":
        break


def exif_size(img):
    """Returns the EXIF-corrected PIL image size as a tuple (width, height)."""
    s = img.size  # (width, height)
    try:
        rotation = dict(img._getexif().items())[orientation]
        if rotation in [6, 8]:  # rotation 270
            s = (s[1], s[0])
    except Exception:
        pass

    return s


def split_rows_simple(file="../data/sm4/out.txt"):  # from utils import *; split_rows_simple()
    """Splits a text file into train, test, and val files based on specified ratios; expects a file path as input."""
    with open(file) as f:
        lines = f.readlines()

    s = Path(file).suffix
    lines = sorted(list(filter(lambda x: len(x) > 0, lines)))
    i, j, k = split_indices(lines, train=0.9, test=0.1, validate=0.0)
    for k, v in {"train": i, "test": j, "val": k}.items():  # key, value pairs
        if v.any():
            new_file = file.replace(s, f"_{k}{s}")
            with open(new_file, "w") as f:
                f.writelines([lines[i] for i in v])


def split_files(out_path, file_name, prefix_path=""):  # split training data
    """Splits file names into separate train, test, and val datasets and writes them to prefixed paths."""
    file_name = list(filter(lambda x: len(x) > 0, file_name))
    file_name = sorted(file_name)
    i, j, k = split_indices(file_name, train=0.9, test=0.1, validate=0.0)
    datasets = {"train": i, "test": j, "val": k}
    for key, item in datasets.items():
        if item.any():
            with open(f"{out_path}_{key}.txt", "a") as file:
                for i in item:
                    file.write(f"{prefix_path}{file_name[i]}\n")


def split_indices(x, train=0.9, test=0.1, validate=0.0, shuffle=True):  # split training data
    """Splits array indices for train, test, and validate datasets according to specified ratios."""
    n = len(x)
    v = np.arange(n)
    if shuffle:
        np.random.shuffle(v)

    i = round(n * train)  # train
    j = round(n * test) + i  # test
    k = round(n * validate) + j  # validate
    return v[:i], v[i:j], v[j:k]  # return indices


def make_dirs(dir="new_dir/"):
    """Creates a directory with subdirectories 'labels' and 'images', removing existing ones."""
    dir = Path(dir)
    if dir.exists():
        shutil.rmtree(dir)  # delete dir
    for p in dir, dir / "labels", dir / "images":
        p.mkdir(parents=True, exist_ok=True)  # make dir
    return dir


def write_data_data(fname="data.data", nc=80):
    """Writes a Darknet-style .data file with dataset and training configuration."""
    lines = [
        f"classes = {nc:g}\n",
        "train =../out/data_train.txt\n",
        "valid =../out/data_test.txt\n",
        "names =../out/data.names\n",
        "backup = backup/\n",
        "eval = coco\n",
    ]

    with open(fname, "a") as f:
        f.writelines(lines)


def image_folder2file(folder="images/"):  # from utils import *; image_folder2file()
    """Generates a txt file listing all images in a specified folder; usage: `image_folder2file('path/to/folder/')`."""
    s = glob.glob(f"{folder}*.*")
    with open(f"{folder[:-1]}.txt", "w") as file:
        for l in s:
            file.write(l + "\n")  # write image list


def add_coco_background(path="../data/sm4/", n=1000):  # from utils import *; add_coco_background()
    """
    Adds COCO dataset background images to a specified folder and lists them in outb.txt; usage:

    `add_coco_background('path/', 1000)`.
    """
    p = f"{path}background"
    if os.path.exists(p):
        shutil.rmtree(p)  # delete output folder
    os.makedirs(p)  # make new output folder

    # copy images
    for image in glob.glob("../coco/images/train2014/*.*")[:n]:
        os.system(f"cp {image} {p}")

    # add to outb.txt and make train, test.txt files
    f = f"{path}out.txt"
    fb = f"{path}outb.txt"
    os.system(f"cp {f} {fb}")
    with open(fb, "a") as file:
        file.writelines(i + "\n" for i in glob.glob(f"{p}/*.*"))
    split_rows_simple(file=fb)


def create_single_class_dataset(path="../data/sm3"):  # from utils import *; create_single_class_dataset('../data/sm3/')
    """Creates a single-class version of an existing dataset in the specified path."""
    os.system(f"mkdir {path}_1cls")


def flatten_recursive_folders(path="../../Downloads/data/sm4/"):  # from utils import *; flatten_recursive_folders()
    """Flattens nested folders in 'path/images' and 'path/json' into single 'images_flat' and 'json_flat'
    directories.
    """
    idir, _jdir = f"{path}images/", f"{path}json/"
    nidir, njdir = Path(f"{path}images_flat/"), Path(f"{path}json_flat/")
    n = 0

    # Create output folders
    for p in [nidir, njdir]:
        if os.path.exists(p):
            shutil.rmtree(p)  # delete output folder
        os.makedirs(p)  # make new output folder

    for parent, dirs, files in os.walk(idir):
        for f in tqdm(files, desc=parent):
            f = Path(f)
            stem, suffix = f.stem, f.suffix
            if suffix.lower()[1:] in img_formats:
                n += 1
                stem_new = f"{n:g}_{stem}"
                image_new = nidir / (stem_new + suffix)  # converts all formats to *.jpg
                json_new = njdir / f"{stem_new}.json"

                image = parent / f
                json = Path(parent.replace("images", "json")) / str(f).replace(suffix, ".json")

                os.system(f"cp '{json}' '{json_new}'")
                os.system(f"cp '{image}' '{image_new}'")
                # cv2.imwrite(str(image_new), cv2.imread(str(image)))

    print(f"Flattening complete: {n:g} jsons and images")


def coco91_to_coco80_class():  # converts 80-index (val2014) to 91-index (paper)
    """Converts COCO 91-class index (paper) to 80-class index (2014 challenge)."""
    return [
        0,
        1,
        2,
        3,
        4,
        5,
        6,
        7,
        8,
        9,
        10,
        None,
        11,
        12,
        13,
        14,
        15,
        16,
        17,
        18,
        19,
        20,
        21,
        22,
        23,
        None,
        24,
        25,
        None,
        None,
        26,
        27,
        28,
        29,
        30,
        31,
        32,
        33,
        34,
        35,
        36,
        37,
        38,
        39,
        None,
        40,
        41,
        42,
        43,
        44,
        45,
        46,
        47,
        48,
        49,
        50,
        51,
        52,
        53,
        54,
        55,
        56,
        57,
        58,
        59,
        None,
        60,
        None,
        None,
        61,
        None,
        62,
        63,
        64,
        65,
        66,
        67,
        68,
        69,
        70,
        71,
        72,
        None,
        73,
        74,
        75,
        76,
        77,
        78,
        79,
        None,
    ]

4.2.2 coco_dataset.py

# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license

import contextlib
import json
from collections import defaultdict

import cv2
import pandas as pd
from PIL import Image

from utils import *


# Convert INFOLKS JSON file into YOLO-format labels ----------------------------
def convert_infolks_json(name, files, img_path):
    """Converts INFOLKS JSON annotations to YOLO-format labels."""
    path = make_dirs()

    # Import json
    data = []
    for file in glob.glob(files):
        with open(file) as f:
            jdata = json.load(f)
            jdata["json_file"] = file
            data.append(jdata)

    # Write images and shapes
    name = path + os.sep + name
    _file_id, file_name, wh, cat = [], [], [], []
    for x in tqdm(data, desc="Files and Shapes"):
        f = glob.glob(img_path + Path(x["json_file"]).stem + ".*")[0]
        file_name.append(f)
        wh.append(exif_size(Image.open(f)))  # (width, height)
        cat.extend(a["classTitle"].lower() for a in x["output"]["objects"])  # categories

        # filename
        with open(name + ".txt", "a") as file:
            file.write(f"{f}\n")

    # Write *.names file
    names = sorted(np.unique(cat))
    # names.pop(names.index('Missing product'))  # remove
    with open(name + ".names", "a") as file:
        [file.write(f"{a}\n") for a in names]

    # Write labels file
    for i, x in enumerate(tqdm(data, desc="Annotations")):
        label_name = Path(file_name[i]).stem + ".txt"

        with open(path + "/labels/" + label_name, "a") as file:
            for a in x["output"]["objects"]:
                # if a['classTitle'] == 'Missing product':
                #    continue  # skip

                category_id = names.index(a["classTitle"].lower())

                # The INFOLKS bounding box format is [x-min, y-min, x-max, y-max]
                box = np.array(a["points"]["exterior"], dtype=np.float32).ravel()
                box[[0, 2]] /= wh[i][0]  # normalize x by width
                box[[1, 3]] /= wh[i][1]  # normalize y by height
                box = [box[[0, 2]].mean(), box[[1, 3]].mean(), box[2] - box[0], box[3] - box[1]]  # xywh
                if (box[2] > 0.0) and (box[3] > 0.0):  # if w > 0 and h > 0
                    file.write("{:g} {:.6f} {:.6f} {:.6f} {:.6f}\n".format(category_id, *box))

    # Split data into train, test, and validate files
    split_files(name, file_name)
    write_data_data(name + ".data", nc=len(names))
    print(f"Done. Output saved to {os.getcwd() + os.sep + path}")


# Convert vott JSON file into YOLO-format labels -------------------------------
def convert_vott_json(name, files, img_path):
    """Converts VoTT JSON files to YOLO-format labels and organizes dataset structure."""
    path = make_dirs()
    name = path + os.sep + name

    # Import json
    data = []
    for file in glob.glob(files):
        with open(file) as f:
            jdata = json.load(f)
            jdata["json_file"] = file
            data.append(jdata)

    # Get all categories
    file_name, wh, cat = [], [], []
    for i, x in enumerate(tqdm(data, desc="Files and Shapes")):
        with contextlib.suppress(Exception):
            cat.extend(a["tags"][0] for a in x["regions"])  # categories

    # Write *.names file
    names = sorted(pd.unique(cat))
    with open(name + ".names", "a") as file:
        [file.write(f"{a}\n") for a in names]

    # Write labels file
    n1, n2 = 0, 0
    missing_images = []
    for i, x in enumerate(tqdm(data, desc="Annotations")):
        f = glob.glob(img_path + x["asset"]["name"] + ".jpg")
        if len(f):
            f = f[0]
            file_name.append(f)
            wh = exif_size(Image.open(f))  # (width, height)

            n1 += 1
            if (len(f) > 0) and (wh[0] > 0) and (wh[1] > 0):
                n2 += 1

                # append filename to list
                with open(name + ".txt", "a") as file:
                    file.write(f"{f}\n")

                # write labelsfile
                label_name = Path(f).stem + ".txt"
                with open(path + "/labels/" + label_name, "a") as file:
                    for a in x["regions"]:
                        category_id = names.index(a["tags"][0])

                        # The INFOLKS bounding box format is [x-min, y-min, x-max, y-max]
                        box = a["boundingBox"]
                        box = np.array([box["left"], box["top"], box["width"], box["height"]]).ravel()
                        box[[0, 2]] /= wh[0]  # normalize x by width
                        box[[1, 3]] /= wh[1]  # normalize y by height
                        box = [box[0] + box[2] / 2, box[1] + box[3] / 2, box[2], box[3]]  # xywh

                        if (box[2] > 0.0) and (box[3] > 0.0):  # if w > 0 and h > 0
                            file.write("{:g} {:.6f} {:.6f} {:.6f} {:.6f}\n".format(category_id, *box))
        else:
            missing_images.append(x["asset"]["name"])

    print(f"Attempted {i:g} json imports, found {n1:g} images, imported {n2:g} annotations successfully")
    if len(missing_images):
        print("WARNING, missing images:", missing_images)

    # Split data into train, test, and validate files
    split_files(name, file_name)
    print(f"Done. Output saved to {os.getcwd() + os.sep + path}")


# Convert ath JSON file into YOLO-format labels --------------------------------
def convert_ath_json(json_dir):  # dir contains json annotations and images
    """Converts ath JSON annotations to YOLO-format labels, resizes images, and organizes data for training."""
    dir = make_dirs()  # output directory

    jsons = []
    for dirpath, dirnames, filenames in os.walk(json_dir):
        jsons.extend(
            os.path.join(dirpath, filename) for filename in [f for f in filenames if f.lower().endswith(".json")]
        )

    # Import json
    n1, n2, n3 = 0, 0, 0
    missing_images, file_name = [], []
    for json_file in sorted(jsons):
        with open(json_file) as f:
            data = json.load(f)

        # # Get classes
        # try:
        #     classes = list(data['_via_attributes']['region']['class']['options'].values())  # classes
        # except:
        #     classes = list(data['_via_attributes']['region']['Class']['options'].values())  # classes

        # # Write *.names file
        # names = pd.unique(classes)  # preserves sort order
        # with open(dir + 'data.names', 'w') as f:
        #     [f.write('%s\n' % a) for a in names]

        # Write labels file
        for x in tqdm(data["_via_img_metadata"].values(), desc=f"Processing {json_file}"):
            image_file = str(Path(json_file).parent / x["filename"])
            f = glob.glob(image_file)  # image file
            if len(f):
                f = f[0]
                file_name.append(f)
                wh = exif_size(Image.open(f))  # (width, height)

                n1 += 1  # all images
                if len(f) > 0 and wh[0] > 0 and wh[1] > 0:
                    label_file = dir + "labels/" + Path(f).stem + ".txt"

                    nlabels = 0
                    try:
                        with open(label_file, "a") as file:  # write labelsfile
                            # try:
                            #     category_id = int(a['region_attributes']['class'])
                            # except:
                            #     category_id = int(a['region_attributes']['Class'])
                            category_id = 0  # single-class

                            for a in x["regions"]:
                                # bounding box format is [x-min, y-min, x-max, y-max]
                                box = a["shape_attributes"]
                                box = np.array(
                                    [box["x"], box["y"], box["width"], box["height"]], dtype=np.float32
                                ).ravel()
                                box[[0, 2]] /= wh[0]  # normalize x by width
                                box[[1, 3]] /= wh[1]  # normalize y by height
                                box = [
                                    box[0] + box[2] / 2,
                                    box[1] + box[3] / 2,
                                    box[2],
                                    box[3],
                                ]  # xywh (left-top to center x-y)

                                if box[2] > 0.0 and box[3] > 0.0:  # if w > 0 and h > 0
                                    file.write("{:g} {:.6f} {:.6f} {:.6f} {:.6f}\n".format(category_id, *box))
                                    n3 += 1
                                    nlabels += 1

                        if nlabels == 0:  # remove non-labelled images from dataset
                            os.system(f"rm {label_file}")
                            # print('no labels for %s' % f)
                            continue  # next file

                        # write image
                        img_size = 4096  # resize to maximum
                        img = cv2.imread(f)  # BGR
                        assert img is not None, "Image Not Found " + f
                        r = img_size / max(img.shape)  # size ratio
                        if r < 1:  # downsize if necessary
                            h, w, _ = img.shape
                            img = cv2.resize(img, (int(w * r), int(h * r)), interpolation=cv2.INTER_AREA)

                        ifile = dir + "images/" + Path(f).name
                        if cv2.imwrite(ifile, img):  # if success append image to list
                            with open(dir + "data.txt", "a") as file:
                                file.write(f"{ifile}\n")
                            n2 += 1  # correct images

                    except Exception:
                        os.system(f"rm {label_file}")
                        print(f"problem with {f}")

            else:
                missing_images.append(image_file)

    nm = len(missing_images)  # number missing
    print(
        f"\nFound {len(jsons):g} JSONs with {n3:g} labels over {n1:g} images. Found {n1 - nm:g} images, labelled {n2:g} images successfully"
    )
    if len(missing_images):
        print("WARNING, missing images:", missing_images)

    # Write *.names file
    names = ["knife"]  # preserves sort order
    with open(dir + "data.names", "w") as f:
        [f.write(f"{a}\n") for a in names]

    # Split data into train, test, and validate files
    split_rows_simple(dir + "data.txt")
    write_data_data(dir + "data.data", nc=1)
    print(f"Done. Output saved to {Path(dir).absolute()}")


def convert_coco_json(json_dir="../coco/annotations/", use_segments=False, cls91to80=False):
    """Converts COCO JSON format to YOLO label format, with options for segments and class mapping."""
    save_dir = make_dirs()  # output directory
    coco80 = coco91_to_coco80_class()

    # Import json
    for json_file in sorted(Path(json_dir).resolve().glob("*.json")):
        fn = Path(save_dir) / "labels" / json_file.stem.replace("instances_", "")  # folder name
        fn.mkdir()
        with open(json_file) as f:
            data = json.load(f)

        # Create image dict
        images = {"{:g}".format(x["id"]): x for x in data["images"]}
        # Create image-annotations dict
        imgToAnns = defaultdict(list)
        for ann in data["annotations"]:
            imgToAnns[ann["image_id"]].append(ann)

        # Write labels file
        for img_id, anns in tqdm(imgToAnns.items(), desc=f"Annotations {json_file}"):
            img = images[f"{img_id:g}"]
            h, w, f = img["height"], img["width"], img["file_name"]

            bboxes = []
            segments = []
            for ann in anns:
                if ann["iscrowd"]:
                    continue
                # The COCO box format is [top left x, top left y, width, height]
                box = np.array(ann["bbox"], dtype=np.float64)
                box[:2] += box[2:] / 2  # xy top-left corner to center
                box[[0, 2]] /= w  # normalize x
                box[[1, 3]] /= h  # normalize y
                if box[2] <= 0 or box[3] <= 0:  # if w <= 0 and h <= 0
                    continue

                cls = coco80[ann["category_id"] - 1] if cls91to80 else ann["category_id"] - 1  # class
                box = [cls] + box.tolist()
                if box not in bboxes:
                    bboxes.append(box)
                # Segments
                if use_segments:
                    if len(ann["segmentation"]) > 1:
                        s = merge_multi_segment(ann["segmentation"])
                        s = (np.concatenate(s, axis=0) / np.array([w, h])).reshape(-1).tolist()
                    else:
                        s = [j for i in ann["segmentation"] for j in i]  # all segments concatenated
                        s = (np.array(s).reshape(-1, 2) / np.array([w, h])).reshape(-1).tolist()
                    s = [cls] + s
                    if s not in segments:
                        segments.append(s)

            # Write
            with open((fn / f).with_suffix(".txt"), "a") as file:
                for i in range(len(bboxes)):
                    line = (*(segments[i] if use_segments else bboxes[i]),)  # cls, box or segments
                    file.write(("%g " * len(line)).rstrip() % line + "\n")


def min_index(arr1, arr2):
    """
    Find a pair of indexes with the shortest distance.

    Args:
        arr1: (N, 2).
        arr2: (M, 2).

    Return:
        a pair of indexes(tuple).
    """
    dis = ((arr1[:, None, :] - arr2[None, :, :]) ** 2).sum(-1)
    return np.unravel_index(np.argmin(dis, axis=None), dis.shape)


def merge_multi_segment(segments):
    """
    Merge multi segments to one list. Find the coordinates with min distance between each segment, then connect these
    coordinates with one thin line to merge all segments into one.

    Args:
        segments(List(List)): original segmentations in coco's json file.
            like [segmentation1, segmentation2,...],
            each segmentation is a list of coordinates.
    """
    s = []
    segments = [np.array(i).reshape(-1, 2) for i in segments]
    idx_list = [[] for _ in range(len(segments))]

    # record the indexes with min distance between each segment
    for i in range(1, len(segments)):
        idx1, idx2 = min_index(segments[i - 1], segments[i])
        idx_list[i - 1].append(idx1)
        idx_list[i].append(idx2)

    # use two round to connect all the segments
    for k in range(2):
        # forward connection
        if k == 0:
            for i, idx in enumerate(idx_list):
                # middle segments have two indexes
                # reverse the index of middle segments
                if len(idx) == 2 and idx[0] > idx[1]:
                    idx = idx[::-1]
                    segments[i] = segments[i][::-1, :]

                segments[i] = np.roll(segments[i], -idx[0], axis=0)
                segments[i] = np.concatenate([segments[i], segments[i][:1]])
                # deal with the first segment and the last one
                if i in [0, len(idx_list) - 1]:
                    s.append(segments[i])
                else:
                    idx = [0, idx[1] - idx[0]]
                    s.append(segments[i][idx[0] : idx[1] + 1])

        else:
            for i in range(len(idx_list) - 1, -1, -1):
                if i not in [0, len(idx_list) - 1]:
                    idx = idx_list[i]
                    nidx = abs(idx[1] - idx[0])
                    s.append(segments[i][nidx:])
    return s


def delete_dsstore(path="../datasets"):
    """Deletes Apple .DS_Store files recursively from a specified directory."""
    from pathlib import Path

    files = list(Path(path).rglob(".DS_store"))
    print(files)
    for f in files:
        f.unlink()


if __name__ == "__main__":
    source = "COCO"

    if source == "COCO":
        convert_coco_json(
            "/workspace/yolo_test/datasets/coco/annotations",  # directory with *.json
            use_segments=True,
            cls91to80=True,
        )

    elif source == "infolks":  # Infolks https://infolks.info/
        convert_infolks_json(name="out", files="../data/sm4/json/*.json", img_path="../data/sm4/images/")

    elif source == "vott":  # VoTT https://github.com/microsoft/VoTT
        convert_vott_json(
            name="data",
            files="../../Downloads/athena_day/20190715/*.json",
            img_path="../../Downloads/athena_day/20190715/",
        )  # images folder

    elif source == "ath":  # ath format
        convert_ath_json(json_dir="../../Downloads/athena/")  # images folder

    # zip results
    # os.system('zip -r ../coco.zip ../coco')

4.2.3 coco.yaml

# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license

# COCO 2017 dataset https://cocodataset.org by Microsoft
# Documentation: https://docs.ultralytics.com/datasets/detect/coco/
# Example usage: yolo train data=coco.yaml
# parent
# ├── ultralytics
# └── datasets
#     └── coco ← downloads here (20.1 GB)

# Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]
path: coco # dataset root dir
train: images/train # train images (relative to 'path') 118287 images
val: images/val # val images (relative to 'path') 5000 images
# test: test-dev2017.txt # 20288 of 40670 images, submit to https://competitions.codalab.org/competitions/20794

# Classes
names:
  0: person
  1: bicycle
  2: car
  3: motorcycle
  4: airplane
  5: bus
  6: train
  7: truck
  8: boat
  9: traffic light
  10: fire hydrant
  11: stop sign
  12: parking meter
  13: bench
  14: bird
  15: cat
  16: dog
  17: horse
  18: sheep
  19: cow
  20: elephant
  21: bear
  22: zebra
  23: giraffe
  24: backpack
  25: umbrella
  26: handbag
  27: tie
  28: suitcase
  29: frisbee
  30: skis
  31: snowboard
  32: sports ball
  33: kite
  34: baseball bat
  35: baseball glove
  36: skateboard
  37: surfboard
  38: tennis racket
  39: bottle
  40: wine glass
  41: cup
  42: fork
  43: knife
  44: spoon
  45: bowl
  46: banana
  47: apple
  48: sandwich
  49: orange
  50: broccoli
  51: carrot
  52: hot dog
  53: pizza
  54: donut
  55: cake
  56: chair
  57: couch
  58: potted plant
  59: bed
  60: dining table
  61: toilet
  62: tv
  63: laptop
  64: mouse
  65: remote
  66: keyboard
  67: cell phone
  68: microwave
  69: oven
  70: toaster
  71: sink
  72: refrigerator
  73: book
  74: clock
  75: vase
  76: scissors
  77: teddy bear
  78: hair drier
  79: toothbrush

数据集目录结构**（coco_dataset.py只是生成了txt文件，需要手动构造以下目录结构）**

datasets/
├── coco/
│   ├── images
|	│   ├── train/
|	│   │   ├── img_001.jpg
|	│   │   ├── img_002.jpg
|	│   │   └── ...
|	│   └── val/
|	│       ├── img_100.jpg
|	│       └── ...
|	└── labels/
|	|   ├── train/
|	|   │   ├── img_001.txt
|	|   │   ├── img_002.txt
|	|   │   └── ...
|	|   └── val/
|	|       ├── img_100.txt
|	|       └── ...
|	└── coco.yaml

5 主代码

└── yolo_test
    ├── bus.jpg
    ├── coco_dataset.py
    ├── datasets
    │   └── coco
    ├── utils.py
    ├── yolo26n.pt
    ├── yolo_predict.py
    ├── yolo_train.py
    └── yolov8n.pt

5.1 训练

from ultralytics import YOLO
import torch
import torch.distributed as dist
import torch_npu
from torch_npu.contrib import transfer_to_npu
torch_npu.npu.set_compile_mode(jit_compile=False)

model = YOLO("yolov8n.pt")
# 多卡训练
# results = model.train(data="./datasets/coco/coco.yaml", epochs=1, device=[0,1,2,3,4,5,6,7], amp=False, batch=256, val=False)
# 单卡训练
results = model.train(data="coco8.yaml", epochs=1, device="0", amp=False, batch=2, val=False)

if __name__ == "__main__":
    if dist.is_initialized():
        dist.destroy_process_group()
    # metrics = model.val(data="coco8.yaml", imgsz=640, batch=16, conf=0.25, iou=0.7, device="0")

启动命令

# 单卡
python train.py
# 多卡
export HCCL_CONNECT_TIMEOUT=6000
torchrun --nproc_per_node=8 yolo_test.py

5.3 推理

from ultralytics import YOLO
from PIL import Image
import torch
import torch_npu
from torch_npu.contrib import transfer_to_npu
torch_npu.npu.set_compile_mode(jit_compile=False)

model = YOLO("yolov8n.pt")
# model.to(device)
results = model.predict("./bus.jpg", device="0")

# Visualize the results
for i, r in enumerate(results):
    im_bgr = r.plot()
    im_rgb = Image.fromarray(im_bgr[..., ::-1])
    r.show()
    r.save(filename=f"results{i}.jpg")

6 报错

6.1运行报错情况

3.1.1RuntimeError: PytorchStreamReader failed reading zip archive: failed finding central directory

若没有下载模型文件，代码运行时会自动下载。下载未完成代码中断了，再次运行时会出现以下报错信息：

RuntimeError: PytorchStreamReader failed reading zip archive: failed finding central directory

请删除原本下载的模型文件，然后重新运行。