rtdetrv2 訓練問題遇到的問題。

pip install torch==2.0.1 torchvision==0.15.2 torchaudio==2.0.2  --index-url https://download.pytorch.org/whl/cu117

1

Please make sure torchvision version >= 0.15.2

發現自己實際裝的是 torchvison==0.15.2+cu117

修改_misc.py中修改為實際版本

if importlib.metadata.version('torchvision') == '0.15.2+cu117':

2

Assertion `index >= -sizes[i] && index < sizes[i] && "index out of bounds"

報錯很多行。

標注的coco序號不對，自己的數據集，分類需要id從0開始，

coco像這種格式。“info”字段也要有，不然也會報錯。

{"info": {"description": "COCO Dataset"},"licenses": [{"name": ""}],"images": [{"id": 1,"file_name": "00002.png","height": 1080,"width": 1920},{"id": 2,"file_name": "00009.png","height": 1080,"width": 1920}],"annotations": [{"id": 1,"image_id": 1,"category_id": 0,"segmentation": [[642.6923076923077,234.23076923076925,1377.3076923076924,234.23076923076925,1377.3076923076924,782.3076923076923,642.6923076923077,782.3076923076923]],"area": 402625.7396449703,"bbox": [642.6923076923077,234.23076923076925,734.6153846153846,548.076923076923],"iscrowd": 0},{"id": 2,"image_id": 2,"category_id": 1,"segmentation": [[490.76923076923083,222.6923076923077,1252.3076923076924,222.6923076923077,1252.3076923076924,784.2307692307692,490.76923076923083,784.2307692307692]],"area": 427633.1360946745,"bbox": [490.76923076923083,222.6923076923077,761.5384615384615,561.5384615384614],"iscrowd": 0}],"categories": [{"id": 0,"name": "ng","supercategory": ""},{"id": 1,"name": "ok","supercategory": ""}]
}

完整轉換腳本，從xml轉為coco。

import os
import json
import xml.etree.ElementTree as ET
from collections import defaultdict
from tqdm import tqdm
import argparse
import shutil
import cv2
import numpy as npdef parse_args():parser = argparse.ArgumentParser(description='Convert Pascal VOC XML annotations to COCO format')parser.add_argument('--xml_dir', type=str, required=True, help='Directory containing XML annotation files')parser.add_argument('--img_dir', type=str, required=True,help='Directory containing corresponding images')parser.add_argument('--output_json', type=str, required=True,help='Output COCO format JSON file path')parser.add_argument('--copy_images', action='store_true',help='Copy images to a new directory structure')parser.add_argument('--output_img_dir', type=str, default='coco_dataset',help='Output directory for images if copying is enabled')return parser.parse_args()def get_image_size(image_path):"""獲取圖像尺寸"""try:img = cv2.imread(image_path)if img is None:raise IOError(f"無法讀取圖像: {image_path}")return img.shape[1], img.shape[0]  # 寬度, 高度except Exception as e:print(f"錯誤: {e}")return 0, 0def convert_xml_to_coco(xml_dir, img_dir, output_json, copy_images=False, output_img_dir=None):# 創建COCO數據結構coco_data = {"info": {"description": "COCO Dataset converted from Pascal VOC XML","version": "1.0","year": 2023,"contributor": "XML to COCO Converter","date_created": "2023-01-01"},"licenses": [{"url": "https://creativecommons.org/licenses/by/4.0/","id": 1,"name": "CC BY 4.0"}],"images": [],"annotations": [],"categories": []}# 處理類別category_dict = {}next_category_id = 0# 處理圖像和標注image_dict = {}next_image_id = 1next_ann_id = 1# 收集所有XML文件xml_files = [f for f in os.listdir(xml_dir) if f.endswith('.xml')]# 創建輸出圖像目錄（如果需要）if copy_images and output_img_dir:os.makedirs(output_img_dir, exist_ok=True)print(f"找到 {len(xml_files)} 個XML文件，開始轉換...")# 處理每個XML文件for xml_file in tqdm(xml_files):xml_path = os.path.join(xml_dir, xml_file)try:# 解析XMLtree = ET.parse(xml_path)root = tree.getroot()# 獲取圖像文件名filename = root.find('filename').textimg_path = os.path.join(img_dir, filename)# 如果圖像不存在，跳過if not os.path.exists(img_path):print(f"警告: 圖像文件不存在 - {img_path}")continue# 獲取圖像尺寸size = root.find('size')if size is not None:width = int(size.find('width').text)height = int(size.find('height').text)else:# 如果XML中沒有尺寸信息，從圖像讀取width, height = get_image_size(img_path)if width == 0 or height == 0:print(f"警告: 無法獲取圖像尺寸 - {img_path}")continue# 如果復制圖像if copy_images and output_img_dir:new_img_path = os.path.join(output_img_dir, filename)shutil.copy2(img_path, new_img_path)# 創建圖像條目if filename not in image_dict:image_entry = {"id": next_image_id,"file_name": filename,"width": width,"height": height,"license": 1,"date_captured": "2023-01-01"}coco_data["images"].append(image_entry)image_dict[filename] = next_image_idnext_image_id += 1image_id = image_dict[filename]# 處理每個對象for obj in root.findall('object'):# 類別處理name = obj.find('name').textif name not in category_dict:category_entry = {"id": next_category_id,"name": name,"supercategory": "object"}coco_data["categories"].append(category_entry)category_dict[name] = next_category_idnext_category_id += 1category_id = category_dict[name]# 邊界框處理bbox = obj.find('bndbox')if bbox is None:continuexmin = float(bbox.find('xmin').text)ymin = float(bbox.find('ymin').text)xmax = float(bbox.find('xmax').text)ymax = float(bbox.find('ymax').text)# 轉換為COCO格式 [x, y, width, height]bbox_width = xmax - xminbbox_height = ymax - ymin# 創建標注條目ann_entry = {"id": next_ann_id,"image_id": image_id,"category_id": category_id,"bbox": [xmin, ymin, bbox_width, bbox_height],"area": bbox_width * bbox_height,"segmentation": [],"iscrowd": 0}coco_data["annotations"].append(ann_entry)next_ann_id += 1except Exception as e:print(f"處理文件 {xml_file} 時出錯: {str(e)}")# 保存為JSON文件with open(output_json, 'w') as f:json.dump(coco_data, f, indent=2)print(f"轉換完成!")print(f"共處理 {len(coco_data['images'])} 張圖像")print(f"共處理 {len(coco_data['annotations'])} 個標注")print(f"共發現 {len(coco_data['categories'])} 個類別")print(f"結果已保存到: {output_json}")# 保存類別映射文件category_map_path = os.path.join(os.path.dirname(output_json), 'category_mapping.txt')with open(category_map_path, 'w') as f:for name, cid in category_dict.items():f.write(f"{name}: {cid}\n")print(f"類別映射已保存到: {category_map_path}")return coco_dataif __name__ == "__main__":args = parse_args()# 運行轉換coco_data = convert_xml_to_coco(args.xml_dir,args.img_dir,args.output_json,args.copy_images,args.output_img_dir)

調用：生成coco的json

python xml_to_coco.py ? ?--xml_dir ?train2017 ? --img_dir ?train2017 ? --output_json ?annotations/instances_train2017.json

python xml_to_coco.py ? ?--xml_dir ?val2017 ? --img_dir ?val2017 ? --output_json ?annotations/instances_val2017.json

數據集結構圖：

然后訓練：

python tools/train.py  --config=configs/rtdetrv2/rtdetrv2_r18vd_120e_coco.yml   --use-amp --seed=0

轉換onnx

python tools/export_onnx.py -c=configs/rtdetrv2/rtdetrv2_r18vd_120e_coco.yml -r last.pth --check

轉換trt，python?版本

python tools/export_trt.py -i model.onnx

或者裝了tensorrt 的用直接命令行。

tensorrt 版本要大于8.5.2，不然有的算子不支持，會報錯。

trtexec --onnx=model.onnx --saveEngine=model.trt

上面python 文件夾whl可以直接pip install?tensorrt-8.6.0-cp39-none-win_amd64.whl

安裝tensorrt python版本。針對直接裝python裝不上的情況。

生成的權重還是挺大的，個人感覺沒有yolo好用。