import colorsys import json import os import matplotlib.pyplot as plt import numpy as np import pylab from keras import backend as K from keras.applications.imagenet_utils import preprocess_input from keras.layers import Input from PIL import Image from tqdm import tqdm from nets.yolo4_tiny import yolo_body, yolo_eval from utils.utils import letterbox_image from yolo import YOLO coco_classes = {'person': 1, 'bicycle': 2, 'car': 3, 'motorbike': 4, 'aeroplane': 5, 'bus': 6, 'train': 7, 'truck': 8, 'boat': 9, 'traffic light': 10, 'fire hydrant': 11, '': 83, 'stop sign': 13, 'parking meter': 14, 'bench': 15, 'bird': 16, 'cat': 17, 'dog': 18, 'horse': 19, 'sheep': 20, 'cow': 21, 'elephant': 22, 'bear': 23, 'zebra': 24, 'giraffe': 25, 'backpack': 27, 'umbrella': 28, 'handbag': 31, 'tie': 32, 'suitcase': 33, 'frisbee': 34, 'skis': 35, 'snowboard': 36, 'sports ball': 37, 'kite': 38, 'baseball bat': 39, 'baseball glove': 40, 'skateboard': 41, 'surfboard': 42, 'tennis racket': 43, 'bottle': 44, 'wine glass': 46, 'cup': 47, 'fork': 48, 'knife': 49, 'spoon': 50, 'bowl': 51, 'banana': 52, 'apple': 53, 'sandwich': 54, 'orange': 55, 'broccoli': 56, 'carrot': 57, 'hot dog': 58, 'pizza': 59, 'donut': 60, 'cake': 61, 'chair': 62, 'sofa': 63, 'pottedplant': 64, 'bed': 65, 'diningtable': 67, 'toilet': 70, 'tvmonitor': 72, 'laptop': 73, 'mouse': 74, 'remote': 75, 'keyboard': 76, 'cell phone': 77, 'microwave': 78, 'oven': 79, 'toaster': 80, 'sink': 81, 'refrigerator': 82, 'book': 84, 'clock': 85, 'vase': 86, 'scissors': 87, 'teddy bear': 88, 'hair drier': 89, 'toothbrush': 90 } class mAP_YOLO(YOLO): #---------------------------------------------------# # 获得所有的分类 #---------------------------------------------------# def generate(self): self.score = 0.01 self.iou = 0.5 model_path = os.path.expanduser(self.model_path) assert model_path.endswith('.h5'), 'Keras model or weights must be a .h5 file.' # 计算anchor数量 num_anchors = len(self.anchors) num_classes = len(self.class_names) # 载入模型，如果原来的模型里已经包括了模型结构则直接载入。 # 否则先构建模型再载入 try: self.yolo_model = load_model(model_path, compile=False) except: self.yolo_model = yolo_body(Input(shape=(None,None,3)), num_anchors//2, num_classes) self.yolo_model.load_weights(self.model_path) else: assert self.yolo_model.layers[-1].output_shape[-1] == \ num_anchors/len(self.yolo_model.output) * (num_classes + 5), \ 'Mismatch between model and given anchor and class sizes' print('{} model, anchors, and classes loaded.'.format(model_path)) # 画框设置不同的颜色 hsv_tuples = [(x / len(self.class_names), 1., 1.) for x in range(len(self.class_names))] self.colors = list(map(lambda x: colorsys.hsv_to_rgb(*x), hsv_tuples)) self.colors = list( map(lambda x: (int(x[0] * 255), int(x[1] * 255), int(x[2] * 255)), self.colors)) # 打乱颜色 np.random.seed(10101) np.random.shuffle(self.colors) np.random.seed(None) self.input_image_shape = K.placeholder(shape=(2, )) #---------------------------------------------------------# # 在yolo_eval函数中，我们会对预测结果进行后处理 # 后处理的内容包括，解码、非极大抑制、门限筛选等 #---------------------------------------------------------# boxes, scores, classes = yolo_eval(self.yolo_model.output, self.anchors, num_classes, self.input_image_shape, max_boxes = self.max_boxes, score_threshold = self.score, iou_threshold = self.iou, letterbox_image=self.letterbox_image) return boxes, scores, classes #---------------------------------------------------# # 检测图片 #---------------------------------------------------# def detect_image(self, image_id, image, results): #---------------------------------------------------------# # 给图像增加灰条，实现不失真的resize #---------------------------------------------------------# if self.letterbox_image: boxed_image = letterbox_image(image, (self.model_image_size[1],self.model_image_size[0])) else: boxed_image = image.convert('RGB') boxed_image = boxed_image.resize((self.model_image_size[1],self.model_image_size[0]), Image.BICUBIC) image_data = np.array(boxed_image, dtype='float32') image_data /= 255. #---------------------------------------------------------# # 添加上batch_size维度 #---------------------------------------------------------# image_data = np.expand_dims(image_data, 0) #---------------------------------------------------------# # 将图像输入网络当中进行预测！ #---------------------------------------------------------# out_boxes, out_scores, out_classes = self.sess.run( [self.boxes, self.scores, self.classes], feed_dict={ self.yolo_model.input: image_data, self.input_image_shape: [image.size[1], image.size[0]], K.learning_phase(): 0}) for i, c in enumerate(out_classes): result = {} predicted_class = self.class_names[c] top, left, bottom, right = out_boxes[i] top = max(0, np.floor(top + 0.5).astype('int32')) left = max(0, np.floor(left + 0.5).astype('int32')) bottom = min(image.size[1], np.floor(bottom + 0.5).astype('int32')) right = min(image.size[0], np.floor(right + 0.5).astype('int32')) result["image_id"] = int(image_id) result["category_id"] = coco_classes[predicted_class] result["bbox"] = [float(left),float(top),float(right-left),float(bottom-top)] result["score"] = float(out_scores[i]) results.append(result) return results yolo = mAP_YOLO() jpg_names = os.listdir("./coco_dataset/val2017") with open("./coco_dataset/eval_results.json","w") as f: results = [] for jpg_name in tqdm(jpg_names): if jpg_name.endswith("jpg"): image_path = "./coco_dataset/val2017/" + jpg_name image = Image.open(image_path) # 开启后在之后计算mAP可以可视化 results = yolo.detect_image(jpg_name.split(".")[0],image,results) json.dump(results,f)