Environment Preparation

Ensure the working directory path contains no Chinese characters. Clone the repository from the official GitHub source using the following command:

git clone https://github.com/ultralytics/yolov5.git

Open the project folder in PyCharm. Verify your CUDA version to ensure compatibility with PyTorch:

nvcc -V

Create a dedicated Conda environment for the project:

conda create -n obj_detection_env python=3.8
conda activate obj_detection_env

Install the required dependencies. To avoid configuration issues, it is recommended to install PyTorch first by selecting the appropriate configuration (CUDA version) from the official PyTorch website, then install the YOLOv5 requirements:

pip install torch torchvision torchaudio
pip install -r requirements.txt

Verify the installation by checking the GPU availability:

import torch

print(f"PyTorch Version: {torch.__version__}")
print(f"CUDA Available: {torch.cuda.is_available()}")
print(f"Device Count: {torch.cuda.device_count()}")
print(f"cuDNN Version: {torch.backends.cudnn.version()}")

Dataset Preparation

Download the coco128 dataset or prepare a custom dataset. The directory structure should follow this format:

• images: Contains subfolders like train2017 for images.
• labels: Contains subfolders like labels2017 for annotation files.

For custom annotation, install labelme:

conda activate obj_detection_env
pip install labelme
python -m labelme

Converting Annotations

LabelMe exports JSON files. Use the script below to convert these JSON files into YOLO-formatted TXT files (normalized center coordinates).

import json
import os

# Mapping class names to IDs
class_mapping = {'face': 0}

def normalize_box(img_dims, bbox):
    img_w, img_h = img_dims
    x1, y1, x2, y2 = bbox
    
    # Calculate center, width, and height
    center_x = ((x1 + x2) / 2) / img_w
    center_y = ((y1 + y2) / 2) / img_h
    width = (x2 - x1) / img_w
    height = (y2 - y1) / img_h
    
    return center_x, center_y, width, height

def process_label_file(json_dir, json_file, output_dir):
    json_path = os.path.join(json_dir, json_file)
    
    with open(json_path, 'r', encoding='utf-8') as f:
        data = json.load(f)
    
    img_width = data['imageWidth']
    img_height = data['imageHeight']
    
    txt_filename = json_file.replace('.json', '.txt')
    txt_path = os.path.join(output_dir, txt_filename)
    
    with open(txt_path, 'w') as txt_out:
        for shape in data['shapes']:
            label = shape['label']
            if label in class_mapping and shape['shape_type'] == 'rectangle':
                points = shape['points']
                # Points are [[x1, y1], [x2, y2]]
                box_coords = (points[0][0], points[0][1], points[1][0], points[1][1])
                
                yolo_bbox = normalize_box((img_width, img_height), box_coords)
                
                line = f"{class_mapping[label]} " + " ".join([str(round(val, 6)) for val in yolo_bbox])
                txt_out.write(line + '\n')

if __name__ == "__main__":
    input_json_folder = 'path/to/labelme/jsons'
    output_txt_folder = 'path/to/yolo/labels/train'
    
    # Ensure output directory exists
    os.makedirs(output_txt_folder, exist_ok=True)
    
    for file in os.listdir(input_json_folder):
        if file.endswith('.json'):
            process_label_file(input_json_folder, file, output_txt_folder)

Configuration Files

Modify the data configuration file (e.g., custom_data.yaml) to point to your dataset paths and define classes:

path: ../datasets/coco128  # dataset root dir
train: images/train2017  # train images
val: images/train2017  # val images

# Classes
nc: 1  # number of classes
names: ['face']  # class names

Update the model configuration file (e.g., yolov5x.yaml) to match the number of classes (nc).

Training Execution

Run the training script. Key parameters to adjust include weights, data config, batch size, and image size:

python train.py --img 640 --batch 16 --epochs 100 --data custom_data.yaml --cfg yolov5x.yaml --weights yolov5x.pt --device 0

Troubleshooting Common Issues

• Dataset Path Errors: Ensure the paths in the YAML file match the actual location of the images and labels folders.
• Font Download Failures: If the script fails to download Arial.Unicode.ttf, manually download the font and place it in the user's Ultralytics AppData folder or the root project directory.
• Virtual Memory (Page File): If encountering WinError 1455, increase the system's virtual memory paging file size.
• Missing Labels: Ensure the converted TXT files are inside the correct labels/train directory specified in the YAML.
• Encoding Errors: If UnicodeDecodeError occurs while reading YAML files, remove any Chinese comments or ensure the file is saved with UTF-8 encoding.