detr输出预测信息

修改coco_eval.py文件

    def prepare_for_coco_detection(self, predictions, output_dir='/home/rendoudou/detr-annotations-master/output'):
        results_by_image_id = {}
        for original_id, prediction in predictions.items():
            if len(prediction) == 0:
                continue

            boxes = prediction["boxes"]
            boxes = convert_to_xywh(boxes).tolist()
            scores = prediction["scores"].tolist()
            labels = prediction["labels"].tolist()

            for k, box in enumerate(boxes):
                category_id = labels[k]
                bbox = [int(x) for x in box]
                score = scores[k]

                result_str = f"{category_id} {' '.join(str(x) for x in bbox)} {score}"

                if original_id not in results_by_image_id:
                    results_by_image_id[original_id] = []

                results_by_image_id[original_id].append(result_str)

        # 转换为对象数组
        coco_results = []
        for image_id, results in results_by_image_id.items():
            for result in results:
                fields = result.split(" ")
                category_id = int(fields[0])
                bbox = [float(x) for x in fields[1:5]]
                score = float(fields[-1])

                coco_result = {
                    "image_id": int(image_id),
                    "category_id": int(category_id),
                    "bbox": bbox,
                    "score": float(score),
                }
                coco_results.append(coco_result)

        # 写入txt文件
        for image_id, results in results_by_image_id.items():
            filename = os.path.join(output_dir, f"{image_id}.txt")
            with open(filename, "w", encoding="utf-8") as f:
                f.writelines("\n".join(results))

        return coco_results
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47