作者：梁云1991
链接：https://mp.weixin.qq.com/s/ezNqSGKxnV1VvShIWV6EUw
本文仅用于学术分享，出于目的，请联系后台作删文处理
另外，这篇文章的内容对同学们熟悉超加工这个库的代码结构有帮助。
江大白后台回复关键词：torchkeras，获取论文笔记本来源代码和气球数据集下载链接。
from ultralytics import YOLO
model = YOLO('yolov8n-seg.pt')
import numpy as np
from PIL import Image
img_path = 'park.jpg'
try:
img = Image.open(img_path)
except Exception as err:
from torchkeras.data import get_example_image
img = get_example_image(img_path)
img.save(img_path)
#可以保存预测结果以及可视化图片
result = model.predict(source=img_path, save=True,
save_txt=True, conf = 0.3)
from pathlib import Path
import ultralytics
from ultralytics.data import utils
yaml_path = str(Path(ultralytics.__file__).parent/'cfg/datasets/coco128-seg.yaml')
class_names = utils.yaml_load(yaml_path)['names']
from torchkeras import plots
boxes = result[0].boxes.data
masks = result[0].masks.data
plots.plot_instance_segmentation(img,boxes,masks,class_names)
yolo_dataset
├── images
│ ├── train
│ │ ├── train0.jpg
│ │ └── train1.jpg
│ ├── val
│ │ ├── val0.jpg
│ │ └── val1.jpg
│ └── test
│ ├── test0.jpg
│ └── test1.jpg
└── labels
├── train
│ ├── train0.txt
│ └── train1.txt
├── val
│ ├── val0.txt
│ └── val1.txt
└── test
├── test0.txt
└── test1.txt
对于实例分割模型，标签文件(如train0.txt)格式如下：
class_id point1(x,y) point2(x,y) point3(x,y) point4(x,y),...
8 0.417781 0.771355 0.440328 0.735397 0.467375 0.658995 0.440328 0.605047 0.387719 0.524159 0.378703 0.443248 0.333625 0.436519 0.371188 0.375841 0.335125 0.364603 0.350156 0.315164 0.320094 0.299439 0.320094 0.256752 0.327609 0.198318 0.357672 0.184836 0.39825 0.155607 0.498937 0.139883 0.470375 0.0724766 0.513953 0.117407 0.553031 0.083715 0.608641 0.115164 0.67175 0.173598 0.704812 0.184836 0.710828 0.211799 0.707828 0.232033 0.718344 0.263481 0.713828 0.308435 0.707828 0.348879 0.691297 0.398318 0.676266 0.416285 0.673266 0.476963 0.641703 0.420794 0.623672 0.510678 0.604125 0.566846 0.560547 0.623037 0.547016 0.676963 0.568063 0.730888 0.607141 0.771355 0.584594 0.811799 0.506453 0.829766 0.411766 0.793832 0.420781 0.769112
注意class_id从0开始，点坐标都是相对图片长宽的相对坐标。
import os,json
from pathlib import Path
from shutil import copyfile
from PIL import Image
from tqdm import tqdm
root_path = './datasets/balloon-seg/'
# 1,构建目录
data_root = Path(root_path)
for tp in ('images','labels'):
for part in ('train','val'):
(data_root/tp/part).mkdir(parents=True, exist_ok=True)
# 2,复制图片文件
train_images = [str(x) for x in Path('balloon/train/').rglob('*.jpg')]
val_images = [str(x) for x in Path('balloon/val/').rglob('*.jpg')]
for src_file in tqdm(train_images):
name = os.path.basename(src_file)
dst_file = root_path+'images/train/'+name
copyfile(src_file,dst_file)
for src_file in tqdm(val_images):
name = os.path.basename(src_file)
dst_file = root_path+'images/val/'+name
copyfile(src_file,dst_file)
# 3,生成标签文件
train_dir = 'balloon/train/'
val_dir = 'balloon/val/'
train_json_file = train_dir + 'via_region_data.json'
val_json_file = val_dir + 'via_region_data.json'
def get_poly(anno):
anno = anno['shape_attributes']
px = anno['all_points_x']
py = anno['all_points_y']
poly = [(x + 0.5, y + 0.5) for x, y in zip(px, py)]
poly = [p for x in poly for p in x]
#box = [np.min(px), np.min(py), np.max(px), np.max(py)]
return poly
def convert_yolo(size,poly):
width,height = size
dh,dw = 1.0/height,1.0/width
poly = [dw*x if i%2==0 else dh*x for i,x in enumerate(poly)]
return poly
def write_yolo_txt(label_path,catids,yolo_polys):
lines = [f'{cls} {' '.join(f'{x:.6f}' for x in poly)}\n'
for cls,poly in zip(catids,yolo_polys)]
with open(label_path, 'w') as fl:
fl.writelines(lines)
def write_labels(data_dir,part):
with open(data_dir + 'via_region_data.json') as f:
info = json.load(f)
info_values = list(info.values())
for info_value in tqdm(info_values):
img_path = data_dir + info_value['filename']
anno_list = list(info_value['regions'].values())
polys = [get_poly(anno) for anno in anno_list]
catids = np.array([0 for x in polys])
size = Image.open(img_path).size
yolo_polys = [convert_yolo(size,poly) for poly in polys]
txt_path = data_root/'labels'/part/info_value['filename'].replace('.jpg','.txt')
write_yolo_txt(txt_path,catids,yolo_polys)
write_labels(train_dir,'train')
write_labels(val_dir,'val')
100%|██████████| 63/63 [00:00<00:00, 2057.52it/s]
100%|██████████| 13/13 [00:00<00:00, 1704.63it/s]
100%|██████████| 61/61 [00:00<00:00, 2880.77it/s]
100%|██████████| 13/13 [00:00<00:00, 2820.50it/s]
from PIL import Image,ImageDraw
import os
from pathlib import Path
from shutil import copyfile
from tqdm import tqdm
import numpy as np
def get_labels_polys(img_path,gt_path):
img = Image.open(img_path)
w,h = img.size
with open(gt_path, 'r') as fl:
lines = [x.rstrip() for x in fl.readlines()]
str_data = [x.split(' ') for x in lines]
relative_polys = [[float(x) for x in arr[1:]] for arr in str_data]
labels = [int(arr[0]) for arr in str_data]
polys = [ [x*w if i%2==0 else x*h for i,x in enumerate(arr)] for arr in relative_polys]
return labels,polys
def plot_polys(image,polys):
image_result = image.copy()
draw = ImageDraw.Draw(image_result)
for poly in polys:
draw.polygon(poly, fill ='cyan', outline ='red')
return image_result
from pathlib import Path
root_path = './datasets/balloon-seg/'
data_root = Path(root_path)
val_imgs = [str(x) for x in (data_root/'images'/'val').rglob('*.jpg') if 'checkpoint' not in str(x)]
img_path = val_imgs[2]
gt_path = img_path.replace('images','labels').replace('.jpg','.txt')
labels,polys = get_labels_polys(img_path,gt_path)
plot_polys(Image.open(img_path),polys)
import ultralytics
print(ultralytics.__file__)
%%writefile balloon-seg.yaml
# 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: balloon-seg # dataset root dir
train: images/train # train images (relative to 'path') 128 images
val: images/val # val images (relative to 'path') 128 images
test: # test images (optional)
# Classes
names:
0: ballon
# Download script/URL (optional)
# download: https://ultralytics.com/assets/coco128.zip
import torch
from torch.utils.data import DataLoader
from ultralytics.cfg import get_cfg
from ultralytics.utils import DEFAULT_CFG,yaml_load
from ultralytics.data.utils import check_cls_dataset, check_det_dataset
from ultralytics.data import build_yolo_dataset,build_dataloader
overrides = {'task':'segment',
'data':'balloon-seg.yaml',
'imgsz':640,
'workers':4
}
cfg = get_cfg(cfg = DEFAULT_CFG,overrides=overrides)
data_info = check_det_dataset(cfg.data)
ds_train = build_yolo_dataset(cfg,img_path=data_info['train'],batch=cfg.batch,
data = data_info,mode='train',rect=False,stride=32)
ds_val = build_yolo_dataset(cfg,img_path=data_info['val'],batch=cfg.batch,data = data_info,
mode='val',rect=False,stride=32)
dl_train = DataLoader(ds_train,batch_size = cfg.batch, num_workers = cfg.workers,
collate_fn = ds_train.collate_fn)
dl_val = DataLoader(ds_val,batch_size = cfg.batch, num_workers = cfg.workers,
collate_fn = ds_val.collate_fn)
from ultralytics.nn.tasks import SegmentationModel
model = SegmentationModel(cfg = 'yolov8n-seg.yaml', ch=3, nc=1)
#weights = torch.hub.load_state_dict_from_url('https://github.com/ultralytics/assets/releases/download/v0.0.0/yolov8n.pt')
weights = torch.load('yolov8n-seg.pt')
model.load(weights['model'])
model.args = cfg
model.nc = data_info['nc'] # attach number of classes to model
model.names = data_info['names']
from ultralytics import YOLO
model = YOLO('yolov8n-seg.pt')
results = model.train(data = 'balloon-seg.yaml', epochs = 50, workers=2)
#自动调参接口
#yolo_model = YOLO('yolov8n-seg.pt')
#yolo_model.tune(data='balloon-seg.yaml', epochs=10,
# iterations=300, optimizer='AdamW',
# plots=True, save=False, val=True)
#测试loss计算过程
for batch in dl_train:
break
for key,value in batch.items():
if isinstance(value,torch.Tensor):
batch[key] = batch[key].cuda()
model = model.cuda()
model.train();
batch['img'] = batch['img'].float()/255.0
loss,_ = model.forward(batch)
loss
张量(89.2061）
from torchkeras import KerasModel
#我们需要修改StepRunner以适应Yolov8的数据集格式
class StepRunner:
def __init__(self, net, loss_fn, accelerator, stage = 'train', metrics_dict = None,
optimizer = None, lr_scheduler = None
):
self.net,self.loss_fn,self.metrics_dict,self.stage = net,loss_fn,metrics_dict,stage
self.optimizer,self.lr_scheduler = optimizer,lr_scheduler
self.accelerator = accelerator
self.net.train()
def __call__(self, batch):
batch['img'] = batch['img'].float()/255
#loss
loss,_ = model.forward(batch)
#backward()
if self.optimizer is not None and self.stage=='train':
self.accelerator.backward(loss)
self.optimizer.step()
if self.lr_scheduler is not None:
self.lr_scheduler.step()
self.optimizer.zero_grad()
all_loss = self.accelerator.gather(loss).sum()
#losses
step_losses = {self.stage+'_loss':all_loss.item()}
#metrics
step_metrics = {}
if self.stage=='train':
if self.optimizer is not None:
step_metrics['lr'] = self.optimizer.state_dict()['param_groups'][0]['lr']
else:
step_metrics['lr'] = 0.0
return step_losses,step_metrics
KerasModel.StepRunner = StepRunner
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)
keras_model = KerasModel(net = model,
loss_fn = None,
optimizer = optimizer)
keras_model.fit(train_data=dl_train,
val_data=dl_val,
epochs = 200,
ckpt_path='checkpoint',
patience=20,
monitor='val_loss',
mode='min',
mixed_precision='no',
plot= True
)
#关闭mosaic增强再训一次
ds_train.close_mosaic(cfg)
keras_model.from_scratch = False
keras_model.fit(train_data=dl_train,
val_data=dl_val,
epochs = 200,
ckpt_path='checkpoint',
patience=20,
monitor='val_loss',
mode='min',
mixed_precision='no',
plot= True
)
from ultralytics import YOLO
keras_model.load_ckpt('checkpoint')
save_dic = dict(model = keras_model.net, train_args =dict(cfg))
torch.save(save_dic, 'best_yolo.pt')
from ultralytics import YOLO
best_model = YOLO(model = 'best_yolo.pt')
metrics = best_model.val(data = cfg.data )
import pandas as pd
df = pd.DataFrame()
df['metric'] = metrics.keys
for i,c in best_model.names.items():
df[c] = metrics.class_result(i)
df
from pathlib import Path
root_path = './datasets/balloon-seg/'
data_root = Path(root_path)
val_imgs = [str(x) for x in (data_root/'images'/'train').rglob('*.jpg') if 'checkpoint' not in str(x)]
img_path = val_imgs[10]
Image.open(img_path)result = best_model(img_path,conf=0.1)
from torchkeras import plots
masks = result[0].masks.data
boxes = result[0].boxes.data
plots.plot_instance_segmentation(Image.open(img_path),boxes,masks,
class_names = ['balloon'],min_score=0.0)
success = best_model.export(format='onnx',dynamic=True)
model = YOLO('best_yolo.onnx',task='segment')
img_path = val_imgs[5]
result = model.predict(img_path,save_txt=True);
from torchkeras import plots
masks = result[0].masks.data
boxes = result[0].boxes.data
plots.plot_instance_segmentation(Image.open(img_path),boxes,masks,
class_names = ['balloon'],min_score=0.0)
import torch
from torch import nn
import onnxruntime
from PIL import Image
from ultralytics.models.yolo.segment.predict import SegmentationPredictor
from ultralytics.yolo.utils.torch_utils import select_device
class OnnxModel(nn.Module):
def __init__(self,weights,
device=torch.device('cpu'),
dnn=False,
data=None,
fp16=False,
fuse=True,
verbose=True):
super().__init__()
w = weights
nn_module = False
onnx = True
pt, jit, xml, engine, coreml, saved_model, pb, tflite, edgetpu, tfjs, paddle, triton = [False]*12
nhwc = False
stride = 32
model, metadata = None, None
cuda = torch.cuda.is_available() and device.type != 'cpu'
names = ['circle']
self.__dict__.update(locals()) # assign all variables to self
providers = ['CUDAExecutionProvider', 'CPUExecutionProvider'] if cuda else [
'CPUExecutionProvider']
self.session = onnxruntime.InferenceSession(w, providers=providers)
self.output_names = [x.name for x in self.session.get_outputs()]
def forward(self, im, augment=False, visualize=False):
im = im.cpu().numpy()
y = self.session.run(self.output_names, {self.session.get_inputs()[0].name: im})
if isinstance(y, (list, tuple)):
return self.from_numpy(y[0]) if len(y) == 1 else [self.from_numpy(x) for x in y]
else:
return self.from_numpy(y)
def from_numpy(self, x):
return torch.tensor(x).to(self.device) if isinstance(x, np.ndarray) else x
def warmup(self,imgsz=(1, 3, 640, 640)):
im = torch.empty(*imgsz, dtype=torch.half if self.fp16 else torch.float,
device=self.device) # input
for _ in range(2 if self.jit else 1): #
self.forward(im) # warmup
class Predictor(SegmentationPredictor):
def setup_model(self, model, verbose=True):
device = select_device(self.args.device, verbose=verbose)
model = model or self.args.model
self.args.half &= device.type != 'cpu'
self.model = OnnxModel(model,
device=device,
dnn=self.args.dnn,
data=self.args.data,
fp16=self.args.half,
verbose=verbose)
self.device = device
self.model.eval()
args = dict(model='best_yolo.onnx')
predictor = Predictor(overrides=args)
result = predictor(source = Image.open(img_path) )
from torchkeras import plots
masks = result[0].masks.data
boxes = result[0].boxes.data
plots.plot_instance_segmentation(Image.open(img_path),boxes,masks,
class_names = ['balloon'],min_score=0.0)
联系客服
