修改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
