#下载源码并解压缩
wget https://github.com/open-mmlab/mmsegmentation/archive/refs/tags/v1.2.1.tar.gz -O mmsegmentation-1.2.1.tar.gz
tar zxvf mmsegmentation-1.2.1.tar.gz
#安装mmcv
pip install mmcv==2.1.0
#安装openmim
pip3 install -U openmim
#安装其它依赖包
cd mmsegmentation-1.2.1
pip3 install -r requirements/runtime.txt
pip3 install -r requirements/tests.txt
pip3 install -r requirements/multimodal.txt
python setup.py develop -i https://pypi.tuna.tsinghua.edu.cn/simple
#下载配置文件和模型文件
mim download mmsegmentation --config segformer_mit-b2_8xb1-160k_cityscapes-1024x1024 --dest .
#测试SegFormer推理单张图片
python demo/image_demo.py demo/demo.png configs/segformer/segformer_mit-b2_8xb1-160k_cityscapes-1024x1024.py segformer_mit-b2_8x1_1024x1024_160k_cityscapes_20211207_134205-6096669a.pth --device cuda:0 --out-file result.jpg
#测试SegFormer推理多张图片
代码仿照demo/image_demo.py修改一下
import os
from glob import glob
from tqdm import tqdm
from argparse import ArgumentParser
from mmengine.model import revert_sync_batchnorm
from mmseg.apis import inference_model, init_model, show_result_pyplot
def main():
parser = ArgumentParser()
parser.add_argument('--input', help='Image file path')
parser.add_argument('--config', help='Config file')
parser.add_argument('--checkpoint', help='Checkpoint file')
parser.add_argument('--output', help='Path to output file')
parser.add_argument(
'--device', default='cuda:0', help='Device used for inference')
parser.add_argument(
'--opacity',
type=float,
default=0.5,
help='Opacity of painted segmentation map. In (0, 1] range.')
parser.add_argument(
'--title', default='result', help='The image identifier.')
args = parser.parse_args()
# build the model from a config file and a checkpoint file
model = init_model(args.config, args.checkpoint, device=args.device)
if args.device == 'cpu':
model = revert_sync_batchnorm(model)
image_files = []
if os.path.exists(args.input):
if os.path.isdir(args.input):
for ext in ['png', 'jpg', 'jpeg', 'bmp']:
files = glob(os.path.join(args.input, '**/*.%s' % (ext)), recursive=True)
if len(files) > 0:
image_files.extend(files)
else:
return
if not os.path.exists(args.output):
os.makedirs(args.output)
for image_path in tqdm(image_files):
(filename, extension) = os.path.splitext(image_path)
filename = os.path.basename(filename)
image_name = "result_" + filename
output_path = os.path.join(args.output, image_name + ".jpg")
# process a single image
result = inference_model(model, image_path)
# show the results
show_result_pyplot(
model,
image_path,
result,
title=args.title,
opacity=args.opacity,
draw_gt=False,
show=False,
out_file=output_path)
if __name__ == '__main__':
main()
#推理测试一个文件夹下的所有图片
python demo/images_dir_demo.py --input demo/images --config configs/segformer/segformer_mit-b2_8xb1-160k_cityscapes-1024x1024.py --checkpoint segformer_mit-b2_8x1_1024x1024_160k_cityscapes_20211207_134205-6096669a.pth --device cuda:0 --output demo/result
统计模型参数量的代码如下:
from mmengine.analysis import get_model_complexity_info
from mmengine.analysis.print_helper import _format_size
from mmseg.apis import inference_model, init_model
from mmengine import Config
import torch
def load_model(config_path, checkpoint_path, resize_height, resize_width, class_num):
cfg = Config.fromfile(config_path)
cfg.crop_size = (resize_height, resize_width)
cfg.data_preprocessor.size = cfg.crop_size
cfg.model.data_preprocessor.size = cfg.crop_size
cfg.test_pipeline = [
dict(type='LoadImageFromFile'),
dict(type='Resize', scale=(resize_width, resize_height), keep_ratio=True),
# add loading annotation after ``Resize`` because ground truth
# does not need to do resize data transform
dict(type='LoadAnnotations'),
dict(type='PackSegInputs')
]
cfg.model.decode_head.num_classes = class_num
print('class_num=', class_num)
model = init_model(cfg, checkpoint_path, 'cuda:0')
return model
if __name__ == '__main__':
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print('==> Building model..')
config_path = 'segformer_mit-b2_8xb1-160k_cityscapes-1024x1024.py'
checkpoint_path = 'segformer_mit-b2_8x1_1024x1024_160k_cityscapes_20211207_134205-6096669a.pth'
# config_path = 'deeplabv3plus_r50b-d8_4xb2-80k_cityscapes-512x1024.py'
# checkpoint_path = 'deeplabv3plus_r50b-d8_512x1024_80k_cityscapes_20201225_213645-a97e4e43.pth'
resize_width = 1024
resize_height = 512
class_num = 19 # Cityscape是19个类别
net = load_model(config_path, checkpoint_path, resize_height, resize_width, class_num)
net = net.to(device=device)
input_shape = (3, resize_height, resize_width)
outputs = get_model_complexity_info(net, input_shape=input_shape, show_table=False, show_arch=False)
flops = _format_size(outputs['flops'])
params = _format_size(outputs['params'])
print("flops:{}".format(flops))
print("params:{}".format(params))