将yolov4预测结果保存为json文件

将yolov4预测结果保存为json文件

如果代码是用mmdetection框架写的，转化部分的代码不需要自己来写，mmdetection自带的脚本可以把预测结果转为json文件。
只需运行python tools/test.py configs/involution/faster_rcnn_red50_neck_fpn_1x_coco.py work_dirs/faster_rcnn_red50_neck_fpn_1x_coco/latest.pth --format-only --options "jsonfile_prefix=./results"即可得到预测结果文件，具体使用方法可以查阅mmdetection的官方文档。
如果没有用到mmdetection框架，可以自己简单修改一下预测部分的代码也是可以将预测结果保存为json文件的。

修改detect_image()函数

参考yolo.py文件中的detect_image()函数，定义save_json函数，返回值不再是带框的图片，而是返回存放框信息的content_json列表，所以在没有检测到物体的时候，返回值不再是原图而是一个空列表。
这里有两种生成json文件的方法，第一种是先读取测试集文件夹下所有的图片名称，然后读取官方给的测试集对应的json文件，如果找到该图片的file_name，则读取image_id，categories_id和预测框bbox以及scores的信息，保存在content_json列表中，在predict.py中调用yolo.py中写好的save_json函数，将结果保存在result.json文件中。

def save_json(self,image,file_name,path_json):
    image_shape = np.array(np.shape(image)[0:2])
        
    if self.letterbox_image:
        crop_img = np.array(letterbox_image(image, (self.model_image_size[1],self.model_image_size[0])))
    else:
        crop_img = image.convert('RGB')
        crop_img = crop_img.resize((self.model_image_size[1],self.model_image_size[0]), Image.BICUBIC)
    photo = np.array(crop_img,dtype = np.float32) / 255.0
    photo = np.transpose(photo, (2, 0, 1))
    
    images = [photo]
    
    with torch.no_grad():
        images = torch.from_numpy(np.asarray(images))
        if self.cuda:
            images = images.cuda()

        outputs = self.net(images)
        output_list = []
        for i in range(3):
            output_list.append(self.yolo_decodes[i](outputs[i]))

        
        output = torch.cat(output_list, 1)
        batch_detections = non_max_suppression(output, len(self.class_names),
                                                conf_thres=self.confidence,
                                                nms_thres=self.iou)
        
        try:
            batch_detections = batch_detections[0].cpu().numpy()
        except:
            return []
        
        top_index = batch_detections[:,4] * batch_detections[:,5] > self.confidence
        top_conf = batch_detections[top_index,4]*batch_detections[top_index,5]
        top_label = np.array(batch_detections[top_index,-1],np.int32)
        top_bboxes = np.array(batch_detections[top_index,:4])
        top_xmin, top_ymin, top_xmax, top_ymax = np.expand_dims(top_bboxes[:,0],-1),np.expand_dims(top_bboxes[:,1],-1),np.expand_dims(top_bboxes[:,2],-1),np.expand_dims(top_bboxes[:,3],-1)

        if self.letterbox_image:
            boxes = yolo_correct_boxes(top_ymin,top_xmin,top_ymax,top_xmax,np.array([self.model_image_size[0],self.model_image_size[1]]),image_shape)
        else:
            top_xmin = top_xmin / self.model_image_size[1] * image_shape[1]
            top_ymin = top_ymin / self.model_image_size[0] * image_shape[0]
            top_xmax = top_xmax / self.model_image_size[1] * image_shape[1]
            top_ymax = top_ymax / self.model_image_size[0] * image_shape[0]
            boxes = np.concatenate([top_ymin,top_xmin,top_ymax,top_xmax], axis=-1)
            
    font = ImageFont.truetype(font='model_data/simhei.ttf',size=np.floor(3e-2 * np.shape(image)[1] + 0.5).astype('int32'))

    thickness = max((np.shape(image)[0] + np.shape(image)[1]) // self.model_image_size[0], 1)
    content_json=[]
    for i, c in enumerate(top_label):
        predicted_class = self.class_names[c]
        score = top_conf[i]

        top, left, bottom, right = boxes[i]
        w=right-left
        h=bottom-top

        y = float(max(0, top))
        x = float(max(0, left))
        w = float(max(0, w))
        h = float(max(0, h))
        score = float(score)
        bbox=[x,y,w,h]

        # print(label, top, left, bottom, right)
        f = open(path_json, 'r', encoding='utf-8') 
        s = f.read()
        rest = json.loads(s)
        for i in rest['images']:
            if i['file_name']!=file_name:
                continue
            if i['file_name']==file_name:
                image_id=i['id']
                cat_id=rest['categories'][c]['id']
                content_dic = {
                    "image_id":image_id,
                    "category_id": cat_id,
                    "bbox": bbox,
                    "score": score
                }
                # print(content_dic)
                content_json.append(content_dic)   
                break
        f.close()  
    return content_json

接着修改predict.py

首先用for i in range(3)输入三张图片测试一下代码有没有问题。

from PIL import Image
import cv2 
from yolo import YOLO
import os
import json
yolo = YOLO()
content=[]
for i in range(3):
    file_name = input('Input image filename:')
    path = "./data/coco/test2017/"
    path_json = "./data/coco/annotations/instances_test2017.json"
    path_save = "./result.json"
    try:
       image = Image.open(path+file_name)
    except:
       print('Open Error! Try again!')
       continue
    else:
       content.extend(yolo.save_json(image,file_name,path_json))
       print(content)
with open(path_save, 'a') as f:
   json.dump(content, f)

没有问题后，用下面的代码将测试集里面的图片全部检测结果保存在result.json文件中。

content=[]
path = "./data/coco/test2017/"
path_json = "./data/coco/annotations/instances_test2017.json"
path_save = "./result.json"
open(path_save, 'w')
for file_name in os.listdir(path):
    image = Image.open(path+file_name)
    info = yolo.save_json(image,file_name,path_json)
    content.extend(info)
with open(path_save, 'a') as f:
    json.dump(content, f)

---

或者用下面的方法，先按照官方文档给出的测试集图片的信息，依次读取图片名，然后取对应的图片进行测试，这样就不用进行上面的判断了。

predict.py

from PIL import Image
import cv2 
from yolo import YOLO
import os
import json
yolo = YOLO()
path = "./data/coco/test2017/"
path_json = "./data/coco/annotations/instances_test2017.json"
path_save = "./result.json"
content=[]
f = open(path_json, 'r', encoding='utf-8') 
s = f.read()
rest = json.loads(s)
for i in rest['images']:
   file_name = i["file_name"]
   image_id = i["id"]
   image = Image.open(path+file_name)
   content.extend(yolo.save_json(image,file_name,image_id,rest['categories']))
f.close()  
with open(path_save, 'a') as f:
   json.dump(content, f)

yolo.py

def save_json(self,image,file_name,image_id,categories):
    image_shape = np.array(np.shape(image)[0:2])
        
    if self.letterbox_image:
        crop_img = np.array(letterbox_image(image, (self.model_image_size[1],self.model_image_size[0])))
    else:
        crop_img = image.convert('RGB')
        crop_img = crop_img.resize((self.model_image_size[1],self.model_image_size[0]), Image.BICUBIC)
    photo = np.array(crop_img,dtype = np.float32) / 255.0
    photo = np.transpose(photo, (2, 0, 1))
    
    images = [photo]
    
    with torch.no_grad():
        images = torch.from_numpy(np.asarray(images))
        if self.cuda:
            images = images.cuda()

        outputs = self.net(images)
        output_list = []
        for i in range(3):
            output_list.append(self.yolo_decodes[i](outputs[i]))

        
        output = torch.cat(output_list, 1)
        batch_detections = non_max_suppression(output, len(self.class_names),
                                                conf_thres=self.confidence,
                                                nms_thres=self.iou)
        
        try:
            batch_detections = batch_detections[0].cpu().numpy()
        except:
            return []
        
        top_index = batch_detections[:,4] * batch_detections[:,5] > self.confidence
        top_conf = batch_detections[top_index,4]*batch_detections[top_index,5]
        top_label = np.array(batch_detections[top_index,-1],np.int32)
        top_bboxes = np.array(batch_detections[top_index,:4])
        top_xmin, top_ymin, top_xmax, top_ymax = np.expand_dims(top_bboxes[:,0],-1),np.expand_dims(top_bboxes[:,1],-1),np.expand_dims(top_bboxes[:,2],-1),np.expand_dims(top_bboxes[:,3],-1)

        if self.letterbox_image:
            boxes = yolo_correct_boxes(top_ymin,top_xmin,top_ymax,top_xmax,np.array([self.model_image_size[0],self.model_image_size[1]]),image_shape)
        else:
            top_xmin = top_xmin / self.model_image_size[1] * image_shape[1]
            top_ymin = top_ymin / self.model_image_size[0] * image_shape[0]
            top_xmax = top_xmax / self.model_image_size[1] * image_shape[1]
            top_ymax = top_ymax / self.model_image_size[0] * image_shape[0]
            boxes = np.concatenate([top_ymin,top_xmin,top_ymax,top_xmax], axis=-1)
            
    font = ImageFont.truetype(font='model_data/simhei.ttf',size=np.floor(3e-2 * np.shape(image)[1] + 0.5).astype('int32'))

    thickness = max((np.shape(image)[0] + np.shape(image)[1]) // self.model_image_size[0], 1)
    content_json=[]
    for i, c in enumerate(top_label):
        predicted_class = self.class_names[c]
        score = top_conf[i]

        top, left, bottom, right = boxes[i]
        w=right-left
        h=bottom-top

        y = float(max(0, top))
        x = float(max(0, left))
        w = float(max(0, w))
        h = float(max(0, h))
        score = float(score)
        bbox=[x,y,w,h]

        # print(label, top, left, bottom, right)
        cat_id=categories[c]['id']
	    content_dic = {
	        "image_id":image_id,
	        "category_id": cat_id,
	        "bbox": bbox,
	        "score": score
	    }
    	content_json.append(content_dic)   
    return content_json