API 구현: utils.py

헬퍼(helper) 함수를 구현한 utils.py 파일을 작성해보겠습니다.

앞에서 살펴본 Faster R-CNN 모델의 예측 결과는 다음과 같은 형태로 반환됩니다:

{
    'boxes': tensor([[x1, y1, x2, y2], ...]),
    'labels': tensor([label1, ...]),
    'scores': tensor([score1, ...]),
}

REST API에서는 각 객체별 결과를 다음과 같이 반환해야 합니다:

{
    "objects": [
    {
        "class": 1,
        "label": "person",
        "score": 0.99,
        "bbox": [x1, y1, x2, y2]
    },
    ...
    ]
}

위와 같이 각 객체 별로 class, label, score, bbox 정보를 반환하도록 헬퍼 함수를 구현해보겠습니다.

# utils.py

# 결과 필터링 함수
def filter_results(outputs, categories, threshold=0.5):
  filtered_results = []

  for label, score, box in zip(outputs['labels'], outputs['scores'], outputs['boxes']):
    if score < threshold:
      continue

    filtered_results.append({
      "class": int(label),
      "label": categories[int(label)],
      "score": float(score),
      "bbox": [float(coord) for coord in box]
    })

  return filtered_results

# 결과 필터링 함수 동작 확인
if __name__ == "__main__":
  sample_outputs = {
    "labels": [1, 1, 2, 3, 4],
    "scores": [0.9, 0.8, 0.7, 0.6, 0.5],
    "boxes": [[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12], [13, 14, 15, 16], [17, 18, 19, 20]]
  }
  sample_categories = {1: "cat", 2: "dog", 3: "bird", 4: "fish"}

  print(filter_results(sample_outputs, sample_categories, 0.75))

[{'class': 1, 'label': 'cat', 'score': 0.9, 'bbox': [1.0, 2.0, 3.0, 4.0]}, {'class': 1, 'label': 'cat', 'score': 0.8, 'bbox': [5.0, 6.0, 7.0, 8.0]}]