使用YOLOv8预测图像时生成相应的json文件

        在使用yolov8对数据进行val验证时，有一个参数save_json,当设置为save_json=True时，模型会将验证时所有关于框的信息存放在一个json文件中。但这个参数只在val验证时可以使用，在predict预测图像时设置save_json=True并不会生成json文件，但我的工作需求是希望在predicr时也能生成图像的json文件，所以我借鉴源码中val里生成json的功能，在源码predicr下也增加了这个功能，然后记录下来。

        可以先看一下val中生成json的代码，在ultralytics/models/yolo/val.py文件中，代码如下： 

    def pred_to_json(self, predn, filename):
        """Serialize YOLO predictions to COCO json format."""
        stem = Path(filename).stem
        image_id = int(stem) if stem.isnumeric() else stem
        box = ops.xyxy2xywh(predn[:, :4])  # xywh
        box[:, :2] -= box[:, 2:] / 2  # xy center to top-left corner
        for p, b in zip(predn.tolist(), box.tolist()):
            self.jdict.append({
                'image_id': image_id,
                'category_id': self.class_map[int(p[5])],
                'bbox': [round(x, 3) for x in b],
                'score': round(p[4], 5)})

        其中最关键的信息是predn， 这是从其他地方传过来的一个参数，具体从哪里生成的因为跟本文关系不大就不再多说，predn信息里包含了关于图片的信息，包括框坐标，图像标签等，但是predn中的信息比较复杂，需要经过处理才能得到有用的信息。

        根据val.py文件，我对同目录下的predict.py文件进行了研究，并对如图所示的pred参数进行了输出。

         输出内容如下：

          

        可以看出，pred是tensor格式，里面是一个二维列表，在后面可以将pred转换为列表格式，当图片经过预测生成两个框时，就会有两个一维数组，有几个框就会有几个一维数组表示每个框的信息。如上图所示，前四个数字表示生成检测框的左上和右下两个点的坐标(x0,y0,x1,y0)。第五个数字表示框的置信度，在检测后生成的图片上也会显示。第六个数字表示框的类别，目前是浮点数，后面可以将其转换为整型。

        有了pred中的内容，接下来就简单了，只需要经过简单的处理，就可以将其中的内容放入json文件，首先可以看到predict.py文件中，类DetectionPredictor继承了类BasePredictor。

 我们可以在目录ultralytics/engine/predictor下找到BasePredictor类，然后在其中加入一个参数self.jdict = [],val中就是这样的，用来存放需要向json中写入的内容。

    # 代码中的表示类别字典，例如：
    label_idx_map = {0:"person", 1:"car", 2:"dog"}

    def pred_to_json(self,pred,img_path):
        """将predict模式下预测结果转换为json格式"""
        # stem表示图片不带后缀的名称
        stem = Path(img_path).stem
        image_id = stem
        # box表示预测的框的坐标
        box = pred[:,:4]
        # 将pred和box转换为列表类型后，依次获取每个框的信息
        for p,b in zip(pred.tolist(),box.tolist()):
            self.jdict.append ({
                'img_id':image_id,
                'img_path':img_path,
                'label':label_idx_map[int(pred[5])], # 表示类别
                 # round（x,4）表示取x小数点后四位
                'bbox':[round(x,4) for x in b],
                'score':round(pred[4],5)
                })
        # 创建json文件的存放路径，self.save_dir是类中自带的默认存放地址
        jsons_path = Path(os.path.join(self.save_dir,'json'))
        if not os.path.exists(jsons_path):
            os.makedirs(jsons_path)
        # 创建json文件，并将self.jdict中的内容写进去
        with open(str(jsons_path / f'{image_id}.json'),'w') as f:
            json.dump(self.jdict,f,ensure_ascii=False)
        # 因为我希望每张图片都可以有一个单独的json文件，所以每次写完一个图片需要将jdict清空
        # 如果想将所有的内容写入一个json文件中(跟val一样),可以删除下面这一行代码
        self.jdict = []

        在源码中我将其写成一个成员方法，可以在predict.py文件下的postprocess方法中使用

    def postprocess(self, preds, img, orig_imgs):
        """Post-processes predictions and returns a list of Results objects."""
        preds = ops.non_max_suppression(preds,
                                        self.args.conf,
                                        self.args.iou,
                                        agnostic=self.args.agnostic_nms,
                                        max_det=self.args.max_det,
                                        classes=self.args.classes)

        results = []
        is_list = isinstance(orig_imgs, list) 
        for i, pred in enumerate(preds):
            orig_img = orig_imgs[i] if is_list else orig_imgs
            if is_list:
                pred[:, :4] = ops.scale_boxes(img.shape[2:], pred[:, :4], orig_img.shape)
            img_path = self.batch[0][i]
            # 代码运行参数，当设置save_json=True时才会使用下面的函数
            if self.args.save_json:
                # 可以看出，我们所写pred_to_json函数需要传入的参数在上面都已经给出，直接使用即可
                self.pred_to_json(pred,img_path)
            results.append(Results(orig_img, path=img_path, names=self.model.names, boxes=pred))
        return results

        到此，将框信息写入json文件的操作就全部完成了。

        这是我写的第一篇博客，不足之处请大家指正，也希望能帮助到大家，并多多交流！