%20(2).svg)
项目简介
本项目展示了如何使用Python和mmdetection框架进行自定义数据集的模型训练。mmdetection是一个基于PyTorch的开源目标检测工具箱，支持多种检测算法和预训练模型。项目主要将LabelMe格式的标注文件转换为COCO格式，并利用转换后的数据集进行模型训练。
项目的主要特性和功能
- 数据转换：通过
labelme2coco.py脚本，可将LabelMe格式的标注文件转换为COCO格式，便于模型训练。 - 图片预处理：使用
resize.py脚本能批量调整图片大小，使其符合模型输入要求。 - 模型训练：借助mmdetection框架提供的工具和配置文件，可对自定义数据集开展模型训练。
- 结果可视化：分析训练日志，绘制准确率和损失值的折线图，还能用训练好的模型进行图像检测。
安装使用步骤
1. 安装依赖
确保已安装Python和Anaconda，按以下步骤配置环境并安装所需库： ```shell conda create -n mmlab python=3.7 -y conda activate mmlab
conda install pytorch torchvision torchaudio pytorch-cuda=11.6 -c pytorch -c nvidia
pip install openmim mim install mmdet ```
2. 数据准备
2.1 图片预处理
使用resize.py脚本批量调整图片大小：
```python
import os
import cv2
img_path = "path_to_your_image_folder" path = os.path.join(img_path) img_list = os.listdir(path)
index = 0 for i in img_list: old_img = cv2.imread(os.path.join(path, i)) new_img = cv2.resize(old_img, (300, 400)) index = index + 1 cv2.imwrite(f"output_folder/{index}.jpg", new_img) ```
2.2 数据标注
使用LabelMe工具进行图片标注，保证输出目录与图片目录一致。
2.3 数据转换
使用labelme2coco.py脚本将LabelMe格式的标注文件转换为COCO格式：
```python
import os
import json
import numpy as np
import glob
import shutil
import cv2
from sklearn.model_selection import train_test_split
np.random.seed(41)
classname_to_id = { "hat": 1, "stalk": 2, }
class Lableme2CoCo: def init(self): self.images = [] self.annotations = [] self.categories = [] self.img_id = 0 self.ann_id = 0
def save_coco_json(self, instance, save_path):
json.dump(instance, open(save_path, 'w', encoding='utf-8'), ensure_ascii=False, indent=1)
def to_coco(self, json_path_list):
self._init_categories()
for json_path in json_path_list:
obj = self.read_jsonfile(json_path)
self.images.append(self._image(obj, json_path))
shapes = obj['shapes']
for shape in shapes:
annotation = self._annotation(shape)
self.annotations.append(annotation)
self.ann_id += 1
self.img_id += 1
instance = {}
instance['info'] = 'spytensor created'
instance['license'] = ['license']
instance['images'] = self.images
instance['annotations'] = self.annotations
instance['categories'] = self.categories
return instance
def _init_categories(self):
for k, v in classname_to_id.items():
category = {}
category['id'] = v
category['name'] = k
self.categories.append(category)
def _image(self, obj, path):
image = {}
from labelme import utils
img_x = utils.img_b64_to_arr(obj['imageData'])
h, w = img_x.shape[:-1]
image['height'] = h
image['width'] = w
image['id'] = self.img_id
image['file_name'] = os.path.basename(path).replace(".json", ".jpg")
return image
def _annotation(self, shape):
label = shape['label']
points = shape['points']
annotation = {}
annotation['id'] = self.ann_id
annotation['image_id'] = self.img_id
annotation['category_id'] = int(classname_to_id[label])
annotation['segmentation'] = [np.asarray(points).flatten().tolist()]
annotation['bbox'] = self._get_box(points)
annotation['iscrowd'] = 0
annotation['area'] = 1.0
return annotation
def read_jsonfile(self, path):
with open(path, "r", encoding='utf-8') as f:
return json.load(f)
def _get_box(self, points):
min_x = min_y = np.inf
max_x = max_y = 0
for x, y in points:
min_x = min(min_x, x)
min_y = min(min_y, y)
max_x = max(max_x, x)
max_y = max(max_y, y)
return [min_x, min_y, max_x - min_x, max_y - min_y]
if name == 'main': labelme_path = "./image" saved_coco_path = "./" print('reading...') if not os.path.exists("%scoco/annotations/" % saved_coco_path): os.makedirs("%scoco/annotations/" % saved_coco_path) if not os.path.exists("%scoco/train2017/" % saved_coco_path): os.makedirs("%scoco/train2017" % saved_coco_path) if not os.path.exists("%scoco/val2017/" % saved_coco_path): os.makedirs("%scoco/val2017" % saved_coco_path) json_list_path = glob.glob(labelme_path + "./*.json") print('json_list_path: ', len(json_list_path)) train_path, val_path = train_test_split(json_list_path, test_size=0.1, train_size=0.9) print("train_n:", len(train_path), 'val_n:', len(val_path))
l2c_train = Lableme2CoCo()
train_instance = l2c_train.to_coco(train_path)
l2c_train.save_coco_json(train_instance, '%scoco/annotations/instances_train2017.json' % saved_coco_path)
for file in train_path:
img_name = file.replace('json', 'jpg')
img_name = img_name[-6:]
temp_img = cv2.imread(os.path.join('image', img_name))
try:
cv2.imwrite("{}coco/train2017/{}".format(saved_coco_path, img_name.replace('png', 'jpg')),temp_img)
except Exception as e:
print(e)
print('Wrong Image:', img_name)
continue
print(img_name + '-->', img_name.replace('png', 'jpg'))
for file in val_path:
img_name = file.replace('json', 'jpg')
img_name = img_name[-6:]
temp_img = cv2.imread(os.path.join('image', img_name))
try:
cv2.imwrite("{}coco/val2017/{}".format(saved_coco_path, img_name.replace('png', 'jpg')), temp_img)
except Exception as e:
print(e)
print('Wrong Image:', img_name)
continue
print(img_name + '-->', img_name.replace('png', 'jpg'))
l2c_val = Lableme2CoCo()
val_instance = l2c_val.to_coco(val_path)
l2c_val.save_coco_json(val_instance, '%scoco/annotations/instances_val2017.json' % saved_coco_path)
```
3. 修改配置文件
依据项目需求，修改mmdetection框架提供的配置文件，以适配自定义数据集。主要修改内容有：
- 修改num_classes为自定义数据集的类别数量。
- 修改img_scale为自定义数据集的图片大小。
- 修改data_root为自定义数据集的根目录。
- 修改ann_file和img_prefix为自定义数据集的标注文件和图片路径。
- 修改samples_per_gpu和workers_per_gpu以适应训练需求。
- 修改optimizer中的学习率以适应训练需求。
4. 训练模型
使用mmdetection框架提供的训练脚本进行模型训练：
shell
python tools/train.py work_dir/mask_rcnn_r50_fpn_1x_coco.py
5. 结果分析
利用提供的分析工具，分析训练日志，绘制准确率和损失值的折线图，并用训练好的模型进行图像检测：
shell
python tools/analysis_tools/analyze_logs.py plot_curve work_dir/20230218_113715.log.json --keys acc
python tools/analysis_tools/analyze_logs.py plot_curve work_dir/20230218_113715.log.json --keys loss_cls loss_bbox loss_mask
python tools/test.py work_dir/mask_rcnn_r50_fpn_1x_coco.py work_dir/epoch_12.pth --show
下载地址
点击下载 【提取码: 4003】【解压密码: www.makuang.net】