1. Pull the code
ONNX > OpenVINO > CoreML > TFLite”>GitHub – ultralytics/ultralytics: NEW – YOLOv8 in PyTorch > ONNX > OpenVINO > CoreML > TFLiteNEW – YOLOv8 in PyTorch > ONNX > OpenVINO > CoreML > TFLite – GitHub – ultralytics/ ultralytics: NEW – YOLOv8 in PyTorch > ONNX > OpenVINO > CoreML > TFLitehttps://github.com/ultralytics/ultralytics .git
2. Build environment
docker pull ultralytics/ultralytics:latest
3. Prepare the data set (labelme annotation, json format)
- Convert json to txt, run labelme2txt.py, and modify the label according to your own data set
import json import os def read_json(json_file): with open(json_file, 'r') as f: load_dict = json.load(f) f.close() return load_dict def mk_file(file_path: str): if os.path.exists(file_path): # If the folder exists, delete the original folder and then recreate it. rmtree(file_path) os.makedirs(file_path) def json2txt(json_path, txt_path): mk_file(txt_path) for json_file in os.listdir(json_path): txt_name = os.path.join(txt_path , json_file[0:-5] + '.txt') txt_file = open(txt_name, 'w') json_file_path = os.path.join(json_path, json_file) json_data = read_json(json_file_path) imageWidth = json_data['imageWidth'] imageHeight = json_data['imageHeight'] for i in range(len(json_data['shapes'])): label = json_data['shapes'][i]['label'] if label == 'car': index = 0 else: index = 1 x1 = json_data['shapes'][i]['points'][0][0] x2 = json_data['shapes'][i]['points'][1][0] y1 = json_data['shapes'][i]['points'][0][1] y2 = json_data['shapes'][i]['points'][1][1] # Compress the label box according to the image size x_center = (x1 + x2) / 2 / imageWidth y_center = (y1 + y2) / 2 / imageHeight bbox_w = (x2 - x1) / imageWidth bbox_h = (y2 - y1) / imageHeight bbox = (x_center, y_center, bbox_w, bbox_h) txt_file.write(str(index) + " " + " ".join([str(a) for a in bbox]) + '\\ ') #print(label) if __name__ == "__main__": json_path = './data/json/' txt_path = './data/txt/' json2txt(json_path, txt_path)
- Divide the data set, run split.py, and modify train_percent to adjust the data set division ratio.
# -*- coding: utf-8 -*- """ Divide the data set into training set, validation set, and test set """ import os import random import shutil #Create a folder to save data def makedir(new_dir): if not os.path.exists(new_dir): os.makedirs(new_dir) def split_data(img_dir, label_dir,save_dir): random.seed(1) # Random seed # 1. Determine the path of the original image data set datasetimg_dir = img_dir # Determine the path of the original label data set datasetlabel_dir = label_dir # 2. Determine the path to save the data set after dividing it. split_dir = save_dir train_dir = os.path.join(split_dir, "train") valid_dir = os.path.join(split_dir, "val") test_dir = os.path.join(split_dir, "test") dir_list = [train_dir, valid_dir, test_dir] image_label = ['images', 'labels'] for i in range(len(dir_list)): for j in range(len(image_label)): makedir(os.path.join(dir_list[i], image_label[j])) # 3. Determine the proportion of dividing the data set into training set, verification set, and test set train_pct = 0.8 valid_pct = 0.2 #test_pct = 0.0 # 4.Divide imgs = os.listdir(datasetimg_dir) # Display all file names in the target folder imgs = list(filter(lambda x: x.endswith('.jpg'), imgs)) # Get all files ending with .png. If the image format is changed, you need to modify it here. random.shuffle(imgs) # shuffle path img_count = len(imgs) # Count the number of images train_point = int(img_count * train_pct) # 0: train_pct valid_point = int(img_count * (train_pct + valid_pct)) # train_pct:valid_pct for i in range(img_count): if i < train_point: # Save the 0-train_point image to the training set out_dir = os.path.join(train_dir, 'images') label_out_dir = os.path.join(train_dir, 'labels') elif i < valid_point: # Save the image of train_point-valid_point to the verification set out_dir = os.path.join(valid_dir, 'images') label_out_dir = os.path.join(valid_dir, 'labels') else: # Save the test_point-end picture to the test set out_dir = os.path.join(test_dir, 'images') label_out_dir = os.path.join(test_dir, 'labels') target_path = os.path.join(out_dir, imgs[i]) #Specify the target saving path src_path = os.path.join(datasetimg_dir, imgs[i]) #Specify the target original image path label_target_path = os.path.join(label_out_dir, imgs[i][0:-4] + '.txt') label_src_path = os.path.join(datasetlabel_dir, imgs[i][0:-4] + '.txt') shutil.copy(src_path, target_path) #Copy pictures shutil.copy(label_src_path, label_target_path) # Copy txt print('train:{}, valid:{}, test:{}'.format(train_point, valid_point - train_point, img_count - valid_point)) if __name__ == "__main__": img_dir = './data/jpg/' label_dir = './data/txt/' save_dir='./data/spilt/' split_data(img_dir, label_dir,save_dir)
4. Modify configuration file
- ultralytics\cfg\datasets\coco128.yaml, modify the path and category name of your own data set image
- ultralytics/cfg/models/v8/yolov8.yaml, modify the number of categories
- ultralytics/cfg/default.yaml Modify the configuration file, path of pre-training weights, number of training times, etc.
5. Run ultralytics/models/yolo/detect/train.py. The latest version cannot be run directly. There is no main function. Add a main function.
def train(cfg=DEFAULT_CFG, use_python=False): """Train and optimize YOLO model given training data and device.""" cfg.model='./yolov8n.pt' model=cfg.model #model = cfg.model or 'yolov8n.pt' data = cfg.data or 'coco128.yaml' # or yolo.ClassificationDataset("mnist") device = cfg.device if cfg.device is not None else '' args = dict(model=model, data=data, device=device) if use_python: from ultralytics import YOLO YOLO(model).train(**args) else: trainer = DetectionTrainer(overrides=args) trainer.train() if __name__ == '__main__': train()