COCOevalについて

import os
import numpy as np
import cv2
import onnxruntime as ort
from pycocotools.coco import COCO
from pycocotools.cocoeval import COCOeval
import json

# 必要なクラスと関数を定義
class ONNXInference:
    def __init__(self, onnx_path, conf_thres=0.25, iou_thres=0.45):
        # ONNX Runtimeセッションの初期化
        self.session = ort.InferenceSession(onnx_path)
        self.input_name = self.session.get_inputs()[0].name
        self.output_names = [output.name for output in self.session.get_outputs()]
        self.conf_threshold = conf_thres
        self.iou_threshold = iou_thres
        
    def preprocess(self, img, input_size=(640, 640)):
        # 前処理
        img = cv2.resize(img, input_size)
        img = img.transpose(2, 0, 1)  # HWC to CHW
        img = img.astype(np.float32) / 255.0
        img = np.expand_dims(img, axis=0)
        return img
        
    def postprocess(self, outputs, original_size):
        # 後処理（YOLOv8の出力を検出ボックスに変換）
        # これはYOLOv8のONNXモデルの出力形式に依存します
        # 以下は簡略化された例です
        predictions = outputs[0]
        
        boxes = []
        scores = []
        class_ids = []
        
        for pred in predictions:
            if pred[4] >= self.conf_threshold:  # 信頼度スコアがしきい値以上
                x1, y1, x2, y2 = pred[0:4]
                score = pred[4]
                class_id = np.argmax(pred[5:])
                
                # 元の画像サイズに合わせてスケーリング
                x1 = x1 * original_size[1] / 640
                y1 = y1 * original_size[0] / 640
                x2 = x2 * original_size[1] / 640
                y2 = y2 * original_size[0] / 640
                
                boxes.append([x1, y1, x2, y2])
                scores.append(score)
                class_ids.append(class_id)
                
        return boxes, scores, class_ids
        
    def infer(self, img):
        original_size = img.shape[:2]
        processed_img = self.preprocess(img)
        outputs = self.session.run(self.output_names, {self.input_name: processed_img})
        boxes, scores, class_ids = self.postprocess(outputs, original_size)
        return boxes, scores, class_ids

# mAPを評価する関数
def evaluate_map(model, annotation_file, image_folder):
    # COCOフォーマットの結果を保存するリスト
    coco_results = []
    
    # アノテーションの読み込み
    coco_gt = COCO(annotation_file)
    image_ids = coco_gt.getImgIds()
    
    # 各画像で推論を実行
    for img_id in image_ids:
        img_info = coco_gt.loadImgs(img_id)[0]
        img_path = os.path.join(image_folder, img_info['file_name'])
        img = cv2.imread(img_path)
        if img is None:
            continue
        
        # 推論
        boxes, scores, class_ids = model.infer(img)
        
        # 結果をCOCOフォーマットに変換
        for box, score, class_id in zip(boxes, scores, class_ids):
            x1, y1, x2, y2 = box
            width = x2 - x1
            height = y2 - y1
            
            # COCOフォーマットの結果
            result = {
                'image_id': img_id,
                'category_id': int(class_id),
                'bbox': [float(x1), float(y1), float(width), float(height)],
                'score': float(score)
            }
            coco_results.append(result)
    
    # 結果をファイルに保存
    with open('detection_results.json', 'w') as f:
        json.dump(coco_results, f)
    
    # COCOevalを使用してmAPを計算
    coco_dt = coco_gt.loadRes('detection_results.json')
    coco_eval = COCOeval(coco_gt, coco_dt, 'bbox')
    coco_eval.evaluate()
    coco_eval.accumulate()
    coco_eval.summarize()
    
    # mAP値を返す
    return coco_eval.stats[0]  # mAP @ IoU=0.5:0.95
    
# 使用例
if __name__ == "__main__":
    # ONNXモデルのパス
    onnx_model_path = "path/to/your/model.onnx"
    
    # アノテーションファイルとデータセットの設定
    annotation_file = "path/to/annotations/instances_val2017.json"
    image_folder = "path/to/images/val2017"
    
    # モデルの初期化
    model = ONNXInference(onnx_model_path)
    
    # mAPの計算
    map_score = evaluate_map(model, annotation_file, image_folder)
    print(f"mAP@0.5:0.95: {map_score}")

pip install pycocotools