要求:实现对汽车的分割。
配置:ubuntu18.1、python3.6.13、pytorch1.7.1、detectron2 0.4、gtx2080ti。
数据集:从kaggle下载
效果如下(迭代1000次):
数据集
1、下载后的文件包括下面这些,这里用到train.zip、train_mask.zip,各5088张,其中train_mask对应的是以二值化掩膜的形式标记的汽车区域;
2、将mask标记的文件转为labelme格式,使用mask_to_labelme.py;因为labelme的json文件中包含图片信息,转换后的json文件可以被labelme.exe打开;
3、将labelme转为coco格式数据集,使用labelme_to_json.py;不想自己转换的可以直接下载;
链接:https://pan.baidu.com/s/15895GbAMBJZ58nqQJe8Osg
提取码:1l92
训练
配置文件跟预训练模型都是detectron2自带的,用默认的就行。
import os
from detectron2.config import get_cfg
from detectron2.engine import DefaultTrainer, default_argument_parser
from detectron2.data.datasets import register_coco_instances
# 数据集下载:https://www.kaggle.com/c/carvana-image-masking-challenge/data
img_path = r"/home/wang/home/wang/dataset/CarvanaImageMaskingChallenge/train"
json_path = r"/home/wang/home/wang/dataset/CarvanaImageMaskingChallenge/Carvana.json"
config_file = r"/home/wang/home/wang/config/COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml"
model_weights = r"/home/wang/home/wang/ModelWeights/model_final_f10217.pkl"
output_dir = r"/home/wang/home/wang/output/Carvana"
register_coco_instances("Carvana", {}, json_path, img_path)
if __name__ == "__main__":
args = default_argument_parser().parse_args()
cfg = get_cfg()
if len(args.config_file) == 0:
args.config_file = config_file
cfg.merge_from_file(args.config_file)
cfg.merge_from_list(args.opts)
cfg.DATASETS.TRAIN = ("Carvana",)
cfg.DATASETS.TEST = () # no metrics implemented for this dataset
cfg.DATALOADER.NUM_WORKERS = 8
cfg.MODEL.WEIGHTS = model_weights
cfg.SOLVER.IMS_PER_BATCH = 4
cfg.SOLVER.BASE_LR = 0.0001
cfg.SOLVER.MAX_ITER = (10000)
cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = (128)
cfg.MODEL.ROI_HEADS.NUM_CLASSES = 1
cfg.OUTPUT_DIR = output_dir
os.makedirs(cfg.OUTPUT_DIR, exist_ok=True)
trainer = DefaultTrainer(cfg)
trainer.resume_or_load(resume=False)
trainer.train()
测试
from detectron2.engine import DefaultPredictor, default_argument_parser
from detectron2.config import get_cfg
from detectron2.data.detection_utils import read_image
from detectron2.utils.visualizer import ColorMode, Visualizer
import os
import matplotlib.pyplot as plt
config_file = r"/home/wang/home/wang/config/COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml"
model_weights = r"/home/wang/home/wang/output/Carvana/model_final.pth"
image_path = r"/home/wang/home/wang/dataset/CarvanaImageMaskingChallenge/predict/0a0e3fb8f782_01.jpg"
out_filename = os.path.join(os.path.dirname(image_path),
os.path.splitext(os.path.basename(image_path))[0] + "_result.jpg")
if __name__ == "__main__":
args = default_argument_parser().parse_args()
cfg = get_cfg()
if len(args.config_file) == 0:
args.config_file = config_file
cfg.merge_from_file(args.config_file)
cfg.merge_from_list(args.opts)
cfg.DATASETS.TEST = ("Carvana",)
cfg.MODEL.WEIGHTS = model_weights
cfg.MODEL.ROI_HEADS.NUM_CLASSES = 1
predictor = DefaultPredictor(cfg)
image_input = read_image(image_path, format="BGR")
image_output = predictor(image_input)
visualizer = Visualizer(image_input, scale=1, instance_mode=ColorMode.IMAGE_BW)
vis_output = visualizer.draw_instance_predictions(image_output["instances"].to("cpu"))
np_output = vis_output.get_image()
plt.subplot(1, 1, 1)
plt.title('PIL')
plt.imshow(np_output)
plt.show()
vis_output.save(out_filename) # 保存