423 lines
19 KiB
Python
Executable File
423 lines
19 KiB
Python
Executable File
import logging
|
|
import os
|
|
from vizer.draw import draw_boxes
|
|
import torch
|
|
import torch.utils.data
|
|
from tqdm import tqdm
|
|
from ssd.config import cfg
|
|
from ssd.data.build import make_data_loader
|
|
from ssd.data.datasets.evaluation import evaluate
|
|
import datetime
|
|
from ssd.utils import dist_util, mkdir
|
|
from ssd.utils.dist_util import synchronize, is_main_process
|
|
# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.
|
|
import logging
|
|
import time
|
|
import os
|
|
import cv2
|
|
import torch
|
|
from tqdm import tqdm
|
|
from ssd.utils.dist_util import is_main_process, get_world_size
|
|
from ssd.utils.dist_util import synchronize
|
|
from ssd.utils.timer import Timer, get_time_str
|
|
|
|
import numpy as np
|
|
from PIL import Image
|
|
|
|
def _accumulate_predictions_from_multiple_gpus(predictions_per_gpu):
|
|
all_predictions = dist_util.all_gather(predictions_per_gpu)
|
|
if not dist_util.is_main_process():
|
|
return
|
|
# merge the list of dicts
|
|
predictions = {}
|
|
for p in all_predictions:
|
|
predictions.update(p)
|
|
# convert a dict where the key is the index in a list
|
|
image_ids = list(sorted(predictions.keys()))
|
|
if len(image_ids) != image_ids[-1] + 1:
|
|
logger = logging.getLogger("SSD.inference")
|
|
logger.warning(
|
|
"Number of images that were gathered from multiple processes is not "
|
|
"a contiguous set. Some images might be missing from the evaluation"
|
|
)
|
|
|
|
# convert to a list
|
|
predictions = [predictions[i] for i in image_ids]
|
|
return predictions
|
|
|
|
|
|
def compute_on_dataset_byHatcherW(model, data_loader, device,timer=None):
|
|
results_dict = {}
|
|
cpu_device = torch.device("cpu")
|
|
now_time = datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')
|
|
filename = os.path.join(
|
|
"/home/baai/voc/HatcherW/SSD/", now_time)
|
|
isExists = os.path.exists(filename)
|
|
if not isExists:
|
|
os.makedirs(filename)
|
|
print(filename + ' 创建成功')
|
|
else:
|
|
print(filename + ' 目录已存在')
|
|
with open(
|
|
"/home/baai/voc/HatcherW/ImageSets/Main/test.txt",
|
|
"r") as f:
|
|
for batch in tqdm(data_loader):
|
|
# 每次进循环先读取一次文件路径
|
|
|
|
|
|
JPEG_NAME = f.readline().strip('\n')
|
|
JPEG_NAME = JPEG_NAME + '.jpg'
|
|
print(JPEG_NAME)
|
|
filename = os.path.join(
|
|
"/home/baai/voc/HatcherW/JPEGImages/",
|
|
JPEG_NAME)
|
|
print(filename)
|
|
####################################
|
|
# src = cv2.imread(filename,flags=1)
|
|
# cv2.namedWindow("yuanlai", cv2.WINDOW_NORMAL)
|
|
# cv2.imshow("yuanlai", src)
|
|
# contrast_brightness_image(src, 0.5, 5) # 第一个1.2为对比度 第二个为亮度数值越大越亮
|
|
|
|
|
|
|
|
########################################
|
|
|
|
|
|
image_forinf = Image.open(filename).convert("RGB") #read image
|
|
image_forinf = np.array(image_forinf)[:, :, [2, 1, 0]] #conver image into array
|
|
# cv2.imshow('测试', image_forinf)
|
|
images, target, image_ids = batch # target in inference does not be used
|
|
|
|
|
|
with torch.no_grad():
|
|
if timer:
|
|
timer.tic()
|
|
output = model(images.to(device))
|
|
if timer:
|
|
if not cfg.MODEL.DEVICE == 'cpu':
|
|
torch.cuda.synchronize()
|
|
timer.toc()
|
|
|
|
output = [o.to(cpu_device) for o in output]
|
|
##########inference 测试,用于测试如何提取数据
|
|
print("output is {}",format(output))
|
|
output_temp = output[0]
|
|
# print("output 类型 is {}".format(type(output[0]['scores'][0])))
|
|
# print("output fenshu is {}".format((output[0]['scores'][0])))
|
|
# print("output biankuang is {}".format((output_temp['boxes'][0,:])))
|
|
# print("output biankuangnum is {}".format((output_temp['boxes'].size())))
|
|
#print("output label is {}".format((output_temp['labels'][0])))
|
|
# print("hhhh")
|
|
bbox_nums = output_temp['boxes'].size()[0] #取出box的个数
|
|
output_scores = output[0]['scores'].numpy()
|
|
output_labels = output[0]['labels'].numpy()
|
|
output_boxes = output[0]['boxes'].numpy()
|
|
|
|
print("output 类型 is {}".format(type(output_scores)))
|
|
result = image_forinf.copy()
|
|
result1 = image_forinf.copy()
|
|
heigh = image_forinf.shape[0]
|
|
kuan = image_forinf.shape[1]
|
|
print(heigh,kuan)
|
|
###visual by hcwang
|
|
for bbox_num in range(bbox_nums): # 画框 需要画bbox_nums个
|
|
left_top = (
|
|
(int)(output_boxes[bbox_num][0] * kuan / 300.0), (int)(output_boxes[bbox_num][1] * heigh / 300.0))
|
|
right_bottom = (
|
|
(int)(output_boxes[bbox_num][2] * kuan / 300.0), (int)(output_boxes[bbox_num][3] * heigh / 300.0))
|
|
cv2.rectangle(result, left_top, right_bottom, (0, 0, 255), 3)
|
|
#cv2.imshow('测试',image_forinf)
|
|
# 存储
|
|
filename = os.path.join("/home/baai/voc/HatcherW/SSD/", now_time ,JPEG_NAME)
|
|
# cv2.imwrite(filename, result)
|
|
for bbox_num in range(bbox_nums):
|
|
output_boxes[bbox_num][0] *=kuan / 300.0
|
|
output_boxes[bbox_num][1] *=heigh / 300.0
|
|
output_boxes[bbox_num][2] *=kuan / 300.0
|
|
output_boxes[bbox_num][3] *=heigh / 300.0
|
|
class_names = ('__background__',
|
|
'aeroplane', 'bicycle', 'bird', 'boat',
|
|
'bottle', 'bus', 'car', 'cat', 'chair',
|
|
'cow', 'diningtable', 'dog', 'horse',
|
|
'motorbike', 'person', 'pottedplant',
|
|
'sheep', 'sofa', 'train', 'tvmonitor')
|
|
drawn_image = draw_boxes(result1, output_boxes, output_labels, output_scores, class_names).astype(np.uint8)
|
|
Image.fromarray(drawn_image).save(filename)
|
|
# cv2.imshow('result',drawn_image) #显示结果
|
|
# cv2.waitKey(0)
|
|
# cv2.destroyAllWindows()
|
|
|
|
print(output)
|
|
results_dict.update(
|
|
{img_id: result for img_id, result in zip(image_ids, output)}
|
|
)
|
|
f.close()
|
|
return results_dict
|
|
def compute_on_dataset(model, data_loader, device):
|
|
results_dict = {}
|
|
for batch in tqdm(data_loader):
|
|
images, targets, image_ids = batch
|
|
cpu_device = torch.device("cpu")
|
|
with torch.no_grad():
|
|
outputs = model(images.to(device))
|
|
|
|
outputs = [o.to(cpu_device) for o in outputs]
|
|
results_dict.update(
|
|
{img_id: result for img_id, result in zip(image_ids, outputs)}
|
|
)
|
|
return results_dict
|
|
|
|
def inference_byHatcherW(model, data_loader, dataset_name, device, output_folder=None, use_cached=False, **kwargs):
|
|
dataset = data_loader.dataset
|
|
logger = logging.getLogger("SSD.inference")
|
|
logger.info("Evaluating {} dataset({} images):".format(dataset_name, len(dataset)))
|
|
total_timer = Timer()
|
|
inference_timer = Timer()
|
|
total_timer.tic()
|
|
|
|
predictions_path = os.path.join(output_folder, 'predictions.pth')
|
|
if use_cached and os.path.exists(predictions_path):
|
|
predictions = torch.load(predictions_path, map_location='cpu')
|
|
else:
|
|
predictions = compute_on_dataset_byHatcherW(model, data_loader, device,inference_timer)
|
|
synchronize()
|
|
predictions = _accumulate_predictions_from_multiple_gpus(predictions)
|
|
if not is_main_process():
|
|
return
|
|
if output_folder:
|
|
torch.save(predictions, predictions_path)
|
|
#return evaluate(dataset=dataset, predictions=predictions, output_dir=output_folder, **kwargs)
|
|
# def inference_byHatcherW(model, data_loader, dataset_name, device, output_folder=None, use_cached=False, **kwargs):
|
|
# dataset = data_loader.dataset
|
|
# logger = logging.getLogger("SSD.inference")
|
|
# logger.info("Evaluating {} dataset({} images):".format(dataset_name, len(dataset)))
|
|
# total_timer = Timer()
|
|
# inference_timer = Timer()
|
|
# total_timer.tic()
|
|
# predictions_path = os.path.join(output_folder, 'predictions.pth')
|
|
# if use_cached and os.path.exists(predictions_path):
|
|
# predictions = torch.load(predictions_path, map_location='cpu')
|
|
# else:
|
|
# predictions = compute_on_dataset(model, data_loader, device,inference_timer)
|
|
# synchronize()
|
|
# predictions = _accumulate_predictions_from_multiple_gpus(predictions)
|
|
# if not is_main_process():
|
|
# return
|
|
# if output_folder:
|
|
# torch.save(predictions, predictions_path)
|
|
# return evaluate(dataset=dataset, predictions=predictions, output_dir=output_folder, **kwargs)
|
|
|
|
def inference(model, data_loader, dataset_name, device, output_folder=None, use_cached=False, **kwargs):
|
|
dataset = data_loader.dataset
|
|
logger = logging.getLogger("SSD.inference")
|
|
logger.info("Evaluating {} dataset({} images):".format(dataset_name, len(dataset)))
|
|
predictions_path = os.path.join(output_folder, 'predictions.pth')
|
|
if use_cached and os.path.exists(predictions_path):
|
|
predictions = torch.load(predictions_path, map_location='cpu')
|
|
else:
|
|
predictions = compute_on_dataset(model, data_loader, device)
|
|
synchronize()
|
|
predictions = _accumulate_predictions_from_multiple_gpus(predictions)
|
|
if not is_main_process():
|
|
return
|
|
if output_folder:
|
|
torch.save(predictions, predictions_path)
|
|
return evaluate(dataset=dataset, predictions=predictions, output_dir=output_folder, **kwargs)
|
|
@torch.no_grad()
|
|
def do_evaluation_byHatcherW(cfg, model, distributed, **kwargs):
|
|
if isinstance(model, torch.nn.parallel.DistributedDataParallel):
|
|
model = model.module
|
|
model.eval()
|
|
device = torch.device(cfg.MODEL.DEVICE)
|
|
data_loaders_val = make_data_loader(cfg, is_train=False, distributed=distributed)
|
|
eval_results = []
|
|
for dataset_name, data_loader in zip(cfg.DATASETS.TEST, data_loaders_val):
|
|
output_folder = os.path.join(cfg.OUTPUT_DIR, "inference", dataset_name)
|
|
if not os.path.exists(output_folder):
|
|
mkdir(output_folder)
|
|
eval_result = inference_byHatcherW(model, data_loader, dataset_name, device, output_folder, **kwargs)
|
|
eval_results.append(eval_result)
|
|
|
|
|
|
def do_evaluation(cfg, model, distributed, **kwargs):
|
|
if isinstance(model, torch.nn.parallel.DistributedDataParallel):
|
|
model = model.module
|
|
model.eval()
|
|
device = torch.device(cfg.MODEL.DEVICE)
|
|
data_loaders_val = make_data_loader(cfg, is_train=False, distributed=distributed)
|
|
eval_results = []
|
|
for dataset_name, data_loader in zip(cfg.DATASETS.TEST, data_loaders_val):
|
|
output_folder = os.path.join(cfg.OUTPUT_DIR, "inference", dataset_name)
|
|
if not os.path.exists(output_folder):
|
|
mkdir(output_folder)
|
|
eval_result = inference(model, data_loader, dataset_name, device, output_folder, **kwargs)
|
|
eval_results.append(eval_result)
|
|
return eval_results
|
|
# def inference_hcw(
|
|
# model,
|
|
# data_loader,
|
|
# dataset_name,
|
|
# iou_types=("bbox",),
|
|
# box_only=False,
|
|
# device="cuda",
|
|
# expected_results=(),
|
|
# expected_results_sigma_tol=4,
|
|
# output_folder=None,
|
|
# ):
|
|
# # convert to a torch.device for efficiency
|
|
# device = torch.device(device) #use GPU
|
|
# num_devices = get_world_size() #get the num of GPU
|
|
# logger = logging.getLogger("maskrcnn_benchmark.inference")
|
|
# dataset = data_loader.dataset #get dataset
|
|
# logger.info("Start evaluation on {} dataset({} images).".format(dataset_name, len(dataset)))
|
|
# total_timer = Timer()
|
|
# inference_timer = Timer()
|
|
# total_timer.tic()
|
|
# predictions = compute_on_dataset(model, data_loader, device, inference_timer)
|
|
# # wait for all processes to complete before measuring the time
|
|
# synchronize()
|
|
# total_time = total_timer.toc()
|
|
# total_time_str = get_time_str(total_time)
|
|
# logger.info(
|
|
# "Total run time: {} ({} s / img per device, on {} devices)".format(
|
|
# total_time_str, total_time * num_devices / len(dataset), num_devices
|
|
# )
|
|
# )
|
|
# total_infer_time = get_time_str(inference_timer.total_time)
|
|
# logger.info(
|
|
# "Model inference time: {} ({} s / img per device, on {} devices)".format(
|
|
# total_infer_time,
|
|
# inference_timer.total_time * num_devices / len(dataset),
|
|
# num_devices,
|
|
# )
|
|
# )
|
|
#
|
|
# predictions = _accumulate_predictions_from_multiple_gpus(predictions)
|
|
# #visual_whc(predictions)
|
|
# if not is_main_process():
|
|
# return
|
|
#
|
|
# if output_folder:
|
|
# torch.save(predictions, os.path.join(output_folder, "predictions.txt"))
|
|
#
|
|
# extra_args = dict(
|
|
# box_only=box_only,
|
|
# iou_types=iou_types,
|
|
# expected_results=expected_results,
|
|
# expected_results_sigma_tol=expected_results_sigma_tol,
|
|
# )
|
|
# def compute_on_dataset_hcw(model, data_loader, device, timer=None):
|
|
# model.eval() #不启用BN和DP
|
|
# results_dict = {} #the result
|
|
# cpu_device = torch.device("cpu")
|
|
# with open(
|
|
# "/home/baai/mask_rcnn/maskrcnn-benchmark/maskrcnn_benchmark/data/datasets/voc/bird/ImageSets/Main/test.txt",
|
|
# "r") as f:
|
|
# for _, batch in enumerate(tqdm(data_loader)):
|
|
# #每次进循环先读取一次文件路径
|
|
# JPEG_NAME = f.readline().strip('\n')
|
|
# JPEG_NAME = JPEG_NAME + '.jpg'
|
|
# print(JPEG_NAME)
|
|
# filename = os.path.join(
|
|
# "/home/baai/mask_rcnn/maskrcnn-benchmark/maskrcnn_benchmark/data/datasets/voc/bird/JPEGImages/",
|
|
# JPEG_NAME)
|
|
# print(filename)
|
|
# image_forinf = Image.open(filename).convert("RGB")
|
|
# image_forinf = np.array(image_forinf)[:, :, [2, 1, 0]]
|
|
# # print("batch is "+batch)
|
|
# images, image_ids = batch
|
|
# # print("image_id is "+image_ids)
|
|
# with torch.no_grad(): #确定是否反向传播求dao
|
|
# if timer:
|
|
# timer.tic()
|
|
# output = model(images.to(device))
|
|
# if timer:
|
|
# if not cfg.MODEL.DEVICE == 'cpu':
|
|
# torch.cuda.synchronize()
|
|
# timer.toc()
|
|
# output = [o.to(cpu_device) for o in output]
|
|
# output_forinf = output[0]
|
|
# print("output caculate success")
|
|
# #将结果rezize
|
|
# height, width = image_forinf.shape[:-1]
|
|
# output_forinf = output_forinf.resize((width, height))
|
|
# print("outout resize success")
|
|
#
|
|
# ################
|
|
# top_predictions = select_top_predictions(output_forinf)
|
|
# print("nms success")
|
|
# result = image_forinf.copy()
|
|
# result = overlay_boxes(result, top_predictions)
|
|
# result = overlay_class_names(result, top_predictions)
|
|
# print("print reslut handle success")
|
|
# ################
|
|
# filename = os.path.join(
|
|
# "/home/baai/mask_rcnn/maskrcnn-benchmark/maskrcnn_benchmark/data/datasets/voc/bird/result1/", JPEG_NAME)
|
|
# cv2.imwrite(filename, result)
|
|
# results_dict.update(
|
|
# {img_id: result for img_id, result in zip(image_ids, output)}
|
|
# )
|
|
# ##为了测试num_bboxex
|
|
# # print("output1:")
|
|
# # temp = results_dict[0].bbox.size()
|
|
# # print(temp[0])
|
|
# # print(temp[1])
|
|
# # print(results_dict[0].bbox.size())
|
|
# #
|
|
# #
|
|
# # print(type(results_dict[0].bbox))
|
|
#
|
|
# # with open(
|
|
# # "/home/baai/mask_rcnn/maskrcnn-benchmark/maskrcnn_benchmark/data/datasets/voc/bird/ImageSets/Main/test.txt",
|
|
# # "r") as f:
|
|
# # JPEG_NAME = f.readline().strip('\n')
|
|
# # # JPEG_NAME.strip('\n')
|
|
# # JPEG_NAME = JPEG_NAME + '.jpg'
|
|
# # print(JPEG_NAME)
|
|
# # f.close()
|
|
# # filename = os.path.join("/home/baai/mask_rcnn/maskrcnn-benchmark/maskrcnn_benchmark/data/datasets/voc/bird/JPEGImages/", JPEG_NAME)
|
|
# # print(filename)
|
|
# # img = cv2.imread(filename)
|
|
# # #cv2.imshow('测试', img)
|
|
# # #cv2.waitKey(0)
|
|
# # #cv2.destroyAllWindows() # dv2.destroyWindow(wname)
|
|
# # results_dict_np = results_dict[0].bbox.numpy()
|
|
# # left_top = ((int)(results_dict_np[0][0]*1920.0/1333.0), (int)(results_dict_np[0][1]*1920.0/1333.0))
|
|
# # right_bottom = ((int)(results_dict_np[0][2]*1920.0/1333.0), (int)(results_dict_np[0][3]*1920.0/1333.0))
|
|
# # cv2.rectangle(img, left_top, right_bottom, (0, 0, 255), 3)
|
|
# # print(results_dict_np[2][1])
|
|
# # filename = os.path.join("/home/baai/mask_rcnn/maskrcnn-benchmark/maskrcnn_benchmark/data/datasets/voc/bird/result/", JPEG_NAME)
|
|
# # cv2.imwrite(filename, img)
|
|
# print('hhhh')
|
|
# print('hhhh')
|
|
# #visual_whc(results_dict)
|
|
# f.close()
|
|
# return results_dict
|
|
def do_evaluation_hcw(cfg, model, distributed, **kwargs):
|
|
if isinstance(model, torch.nn.parallel.DistributedDataParallel):
|
|
model = model.module
|
|
model.eval()
|
|
device = torch.device(cfg.MODEL.DEVICE)
|
|
data_loaders_vals = make_data_loader(cfg, is_train=False, distributed=distributed)
|
|
eval_results = []
|
|
output_folders = [None] * len(cfg.DATASETS.TEST)
|
|
if cfg.OUTPUT_DIR:
|
|
for idx, dataset_name in enumerate(cfg.DATASETS.TEST):
|
|
output_folder = os.path.join(cfg.OUTPUT_DIR, "inference", dataset_name)
|
|
mkdir(output_folder)
|
|
output_folders[idx] = output_folder
|
|
for output_folder, data_loaders_val in zip(output_folders,cfg.DATASETS.TEST, data_loaders_vals):
|
|
eval_result = inference(model, data_loaders_val, dataset_name, device, output_folder, **kwargs)
|
|
eval_results.append(eval_result)
|
|
return eval_results
|
|
|
|
|
|
def contrast_brightness_image(src1, a, g):
|
|
h, w, ch = src1.shape # 获取shape的数值，height和width、通道
|
|
|
|
# 新建全零图片数组src2,将height和width，类型设置为原图片的通道类型(色素全为零，输出为全黑图片)
|
|
src2 = np.zeros([h, w, ch], src1.dtype)
|
|
dst = cv2.addWeighted(src1, a, src2, 1 - a, g) # addWeighted函数说明如下
|
|
cv2.imshow("con-bri-demo", dst)
|
|
|