117 lines
5.1 KiB
Python
117 lines
5.1 KiB
Python
|
import json
|
|||
|
|
import os
|
|||
|
|
|
|||
|
|
import numpy as np
|
|||
|
|
import torch
|
|||
|
|
from PIL import Image
|
|||
|
|
from pycocotools.coco import COCO
|
|||
|
|
from pycocotools.cocoeval import COCOeval
|
|||
|
|
from tqdm import tqdm
|
|||
|
|
|
|||
|
|
from utils.utils import cvtColor, preprocess_input, resize_image
|
|||
|
|
from yolo import YOLO
|
|||
|
|
|
|||
|
|
# ---------------------------------------------------------------------------#
|
|||
|
|
# map_mode用于指定该文件运行时计算的内容
|
|||
|
|
# map_mode为0代表整个map计算流程，包括获得预测结果、计算map。
|
|||
|
|
# map_mode为1代表仅仅获得预测结果。
|
|||
|
|
# map_mode为2代表仅仅获得计算map。
|
|||
|
|
# ---------------------------------------------------------------------------#
|
|||
|
|
map_mode = 0
|
|||
|
|
# -------------------------------------------------------#
|
|||
|
|
# 指向了验证集标签与图片路径
|
|||
|
|
# -------------------------------------------------------#
|
|||
|
|
cocoGt_path = 'coco_dataset/annotations/instances_val2017.json'
|
|||
|
|
dataset_img_path = 'coco_dataset/val2017'
|
|||
|
|
# -------------------------------------------------------#
|
|||
|
|
# 结果输出的文件夹，默认为map_out
|
|||
|
|
# -------------------------------------------------------#
|
|||
|
|
temp_save_path = 'map_out/coco_eval'
|
|||
|
|
|
|||
|
|
|
|||
|
|
class mAP_YOLO(YOLO):
|
|||
|
|
# ---------------------------------------------------#
|
|||
|
|
# 检测图片
|
|||
|
|
# ---------------------------------------------------#
|
|||
|
|
def detect_image(self, image_id, image, results):
|
|||
|
|
# ---------------------------------------------------#
|
|||
|
|
# 计算输入图片的高和宽
|
|||
|
|
# ---------------------------------------------------#
|
|||
|
|
image_shape = np.array(np.shape(image)[0:2])
|
|||
|
|
# ---------------------------------------------------------#
|
|||
|
|
# 在这里将图像转换成RGB图像，防止灰度图在预测时报错。
|
|||
|
|
# 代码仅仅支持RGB图像的预测，所有其它类型的图像都会转化成RGB
|
|||
|
|
# ---------------------------------------------------------#
|
|||
|
|
image = cvtColor(image)
|
|||
|
|
# ---------------------------------------------------------#
|
|||
|
|
# 给图像增加灰条，实现不失真的resize
|
|||
|
|
# 也可以直接resize进行识别
|
|||
|
|
# ---------------------------------------------------------#
|
|||
|
|
image_data = resize_image(image, (self.input_shape[1], self.input_shape[0]), self.letterbox_image)
|
|||
|
|
# ---------------------------------------------------------#
|
|||
|
|
# 添加上batch_size维度
|
|||
|
|
# ---------------------------------------------------------#
|
|||
|
|
image_data = np.expand_dims(np.transpose(preprocess_input(np.array(image_data, dtype='float32')), (2, 0, 1)), 0)
|
|||
|
|
|
|||
|
|
with torch.no_grad():
|
|||
|
|
images = torch.from_numpy(image_data)
|
|||
|
|
if self.cuda:
|
|||
|
|
images = images.cuda()
|
|||
|
|
# ---------------------------------------------------------#
|
|||
|
|
# 将图像输入网络当中进行预测！
|
|||
|
|
# ---------------------------------------------------------#
|
|||
|
|
outputs = self.net(images)
|
|||
|
|
outputs = self.bbox_util.decode_box(outputs)
|
|||
|
|
# ---------------------------------------------------------#
|
|||
|
|
# 将预测框进行堆叠，然后进行非极大抑制
|
|||
|
|
# ---------------------------------------------------------#
|
|||
|
|
outputs = self.bbox_util.non_max_suppression(torch.cat(outputs, 1), self.num_classes, self.input_shape,
|
|||
|
|
image_shape, self.letterbox_image, conf_thres=self.confidence,
|
|||
|
|
nms_thres=self.nms_iou)
|
|||
|
|
|
|||
|
|
if outputs[0] is None:
|
|||
|
|
return results
|
|||
|
|
|
|||
|
|
top_label = np.array(outputs[0][:, 6], dtype='int32')
|
|||
|
|
top_conf = outputs[0][:, 4] * outputs[0][:, 5]
|
|||
|
|
top_boxes = outputs[0][:, :4]
|
|||
|
|
|
|||
|
|
for i, c in enumerate(top_label):
|
|||
|
|
result = {}
|
|||
|
|
top, left, bottom, right = top_boxes[i]
|
|||
|
|
|
|||
|
|
result["image_id"] = int(image_id)
|
|||
|
|
result["category_id"] = clsid2catid[c]
|
|||
|
|
result["bbox"] = [float(left), float(top), float(right - left), float(bottom - top)]
|
|||
|
|
result["score"] = float(top_conf[i])
|
|||
|
|
results.append(result)
|
|||
|
|
return results
|
|||
|
|
|
|||
|
|
|
|||
|
|
if __name__ == "__main__":
|
|||
|
|
if not os.path.exists(temp_save_path):
|
|||
|
|
os.makedirs(temp_save_path)
|
|||
|
|
|
|||
|
|
cocoGt = COCO(cocoGt_path)
|
|||
|
|
ids = list(cocoGt.imgToAnns.keys())
|
|||
|
|
clsid2catid = cocoGt.getCatIds()
|
|||
|
|
|
|||
|
|
if map_mode == 0 or map_mode == 1:
|
|||
|
|
yolo = mAP_YOLO(confidence=0.001, nms_iou=0.65)
|
|||
|
|
|
|||
|
|
with open(os.path.join(temp_save_path, 'eval_results.json'), "w") as f:
|
|||
|
|
results = []
|
|||
|
|
for image_id in tqdm(ids):
|
|||
|
|
image_path = os.path.join(dataset_img_path, cocoGt.loadImgs(image_id)[0]['file_name'])
|
|||
|
|
image = Image.open(image_path)
|
|||
|
|
results = yolo.detect_image(image_id, image, results)
|
|||
|
|
json.dump(results, f)
|
|||
|
|
|
|||
|
|
if map_mode == 0 or map_mode == 2:
|
|||
|
|
cocoDt = cocoGt.loadRes(os.path.join(temp_save_path, 'eval_results.json'))
|
|||
|
|
cocoEval = COCOeval(cocoGt, cocoDt, 'bbox')
|
|||
|
|
cocoEval.evaluate()
|
|||
|
|
cocoEval.accumulate()
|
|||
|
|
cocoEval.summarize()
|
|||
|
|
print("Get map done.")
|