detectron2 测试训练好的模型并画出结果图

最新推荐文章于 2024-08-30 20:56:18 发布

竹子熊猫

最新推荐文章于 2024-08-30 20:56:18 发布

阅读量2.9k

点赞数 4

分类专栏： deep learning python

本文链接：https://blog.csdn.net/summermaoz/article/details/118277770

版权

python 同时被 2 个专栏收录

63 篇文章 1 订阅

订阅专栏

deep learning

35 篇文章 0 订阅

订阅专栏

import numpy as np
import cv2
import os
from PIL import Image
#from matplotlib import pyplot
import matplotlib.pyplot as plt
import random
#from google.colab.patches import cv2_imshow

import detectron2
from detectron2.utils.logger import setup_logger
setup_logger()

from detectron2.engine import DefaultPredictor
from detectron2.config import get_cfg
from detectron2.utils.visualizer import Visualizer
from detectron2.data import MetadataCatalog
from bmaskrcnn import add_boundary_preserving_config

im_folder= './datasets/glass/val/JPEGImages'
save_folder = './result_struc'

for im_file in os.listdir(im_folder):
	im = cv2.imread(os.path.join(im_folder, im_file))

	save_result_path = os.path.join(save_folder, im_file)

	height = im.shape[0]
	width = im.shape[1]
	dpi = 500

	cfg = get_cfg()
	add_boundary_preserving_config(cfg)
	cfg.merge_from_file('./configs/bmask_rcnn_R_101_FPN_3x.yaml')
	cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5  #模型阈值
	cfg.MODEL.WEIGHTS = 'output/bmask_rcnn_r101_3x-structural/model_final.pth'
	predictor = DefaultPredictor(cfg)
	outputs = predictor(im)

	pred_classes = outputs["instances"].pred_classes
	pred_boxes = outputs["instances"].pred_boxes

	#在原图上画出检测结果
	v = Visualizer(im[:, :, ::-1], MetadataCatalog.get(cfg.DATASETS.TRAIN[0]), scale=1)
	v = v.draw_instance_predictions(outputs["instances"].to("cpu"))

	plt.figure(figsize=(width/dpi, height/dpi), dpi=dpi)
	plt.axis('off')
	plt.subplots_adjust(top = 1, bottom = 0, right = 1, left = 0, hspace = 0, wspace = 0)
	plt.imshow(v.get_image())
	plt.savefig(save_result_path) #保存结果