[Label Studio] Ultralytics YOLOv8로 라벨링 작업 자동화 해보기

개요

YOLO와 같은 모델을 학습하기 위해 데이터셋을 준비하는 과정에서 라벨 수 백, 수 천개를 치는 과정은 너무 귀찮고 손이 아프다.

 

이러한 번거로운 과정을 오토라벨링을 통해 조금이라도 단축시킬 수 있다면 어떨까?

 

https://github.com/seblful/label-studio-yolov8-backend

 

그 방법은 바로 위 yolov8 backend 코드를 docker compose를 통해 실행시켜서 Label Studio에 사용하는 것이다.

 

준비물

• 당연하겠지만 Label Studio 서버
• Git
• Docker

설치

1. git clone

git clone https://github.com/seblful/label-studio-yolov8-backend.git
cd label-studio-yolov8-backend

 

2. .env 파일 작성

LABEL_STUDIO_URL=<IPv4 Address> (check your ipconfig)
LABEL_STUDIO_API_KEY=<Label Studio API token>
TASK_TYPE=<segmentation> or <detection>

말 그대로 .env 라는 파일을 만들어서 위 포맷에 맞게 작성하면 된다.

 

LABEL_STUDIO_URL : 사용중인 Label Studio 서버의 IP

LABEL_STUDIO_API_KEY : 사용중인 Label Studio에 생성한 계정의 Access Token

TASK_TYPE : "detection"

Access Token은 Label Studio의 Account & Settings에서 확인할 수 있다.

 

3. model.py 파일 수정

        self.model = ultralytics.YOLO(
            os.path.join(os.getenv("MODEL_DIR"), "best.pt"))

 

models 폴더를 만들고 Ultralytics의 YOLOv8 모델 파일을 넣거나 원하는 모델 파일을 넣는다.

참고로 YOLOv8이 아니더라도 YOLOv10, YOLO11 파일을 넣었을 때도 잘 작동하는것으로 보아 Ultralytics에서 제공하는 YOLO 모델 파일은 다 호환이 가능한 것으로 보인다.

        predictions = [{"result": results,
                       "score": score / counter,
                        "model_version": self.model_version}]

        predictions = [{"result": results,
                       "score": score / counter if counter != 0 else 0.0,
                        "model_version": self.model_version}]

 

"score"의 score / counter 부분을 아래와 같이 변경하고 저장한다.

 

4. (옵션) Dockerfile 수정

# FROM pytorch/pytorch:2.3.1-cuda12.1-cudnn8-runtime
FROM python:3.10-slim

Jetson Jetpack을 사용중인 경우 Dockerfile의 첫 번째 줄을 python:3.10-slim으로 변경해준다.

 

5. YOLO 추론 서버 실행

docker compose up

[+] Running 1/0
 ✔ Container yolo  Created                                                                                                                                                                                                                                                                                     0.0s 
Attaching to yolo
yolo  | [2024-11-14 06:12:07 +0000] [1] [INFO] Starting gunicorn 22.0.0
yolo  | [2024-11-14 06:12:07 +0000] [1] [INFO] Listening at: http://0.0.0.0:9090 (1)
yolo  | [2024-11-14 06:12:07 +0000] [1] [INFO] Using worker: gthread
yolo  | [2024-11-14 06:12:07 +0000] [8] [INFO] Booting worker with pid: 8

 

6. 사용하고자 하는 Project에서 Model 탭 선택 후 Connect Model 클릭

 

7. 정보 입력 후 Validate and Save 버튼 클릭

Name : 아무렇게나 입력

Backend URL : docker compose로 실행했던 9090 포트를 사용하는 IP 주소 입력

 

 

정상적으로 등록되었다면 Connected 라고 뜨고 확인했다면 Save 버튼을 눌러서 저장한다.

 

9. 추론 결과 받아오기

 

모든 Task를 선택하고 Retrieve Predictions를 클릭하면 Backend API를 통해 라벨링 작업을 자동화할 수 있다.

 

이렇게 끝나면 좋겠지만 Retrieve Predictions 버튼을 통해 한 번에 라벨링할 수 있는건 50개인가 100개까지 밖에 되지 않는다.

 

즉, Label Studio의 내장 API를 이용해서 모든 Task에 Backend API의 결과를 삽입하는 코드를 작성해야만 한다.

오토라벨링 코드

1. 오토라벨링 코드 작성

import requests
import json

# Label Studio URL 및 API 토큰
base_url = "http://<IP 주소>:8080"
project_id = "<프로젝트 ID>"
tasks_url = f"{base_url}/api/projects/{project_id}/tasks"
actions_url = f"{base_url}/api/dm/actions?id=retrieve_tasks_predictions&project={project_id}"
api_token = "<API 토큰>"

# 헤더 설정
headers = {
    "Authorization": f"Token {api_token}",
    "Content-Type": "application/json"
}

# 페이지네이션 처리 함수
def get_all_tasks(_filter=None):
    tasks = []
    page = 1

    while True:
    # for i in range(1):
        params = {"page": page}
        response = requests.get(tasks_url, headers=headers, params=params)
        if response.status_code != 200:
            break
        data_list = response.json()

        if not data_list:
            break
        tasks.extend(data_list)
        page +=1
    return tasks

def process():
    tasks = get_all_tasks()
    if tasks:
        print(f"Total tasks retrieved: {len(tasks)}")
    else:
        print("Failed to retrieve tasks")
    
    max_task = len(tasks)
    current_count = 1
    for i, task in enumerate(tasks):
        task_id = task['id']
        data = {
            "selectedItems": {
                "all": False,
                "included": [task_id]
            },
            "project": project_id
        }
        response = requests.post(actions_url, headers=headers, data=json.dumps(data))
        if response.status_code == 200:
            print(f"Successfully retrieved task prediction for Task ID: {task_id} {current_count}/{max_task}")
            current_count += 1
        else:
            print(f"Failed to retrieve task prediction for Task ID: {task_id}")
            print(f"Status Code: {response.status_code}")
            print(response.text)

if __name__ == '__main__':
    process()

 

base_url : Label Studio 주소

project_id : 아무 프로젝트를 클릭하고 주소를 확인 -> http://<주소>:8080/projects/1/ <-- 여기서 "1" 이게 프로젝트 ID

api_token : 아까 입력했던 Access Token 입력

 

2. 오토라벨링 코드 실행

Total tasks retrieved: 162
Successfully retrieved task prediction for Task ID: 2361 1/162
Successfully retrieved task prediction for Task ID: 2363 2/162
Successfully retrieved task prediction for Task ID: 2362 3/162
Successfully retrieved task prediction for Task ID: 2364 4/162
Successfully retrieved task prediction for Task ID: 2365 5/162
Successfully retrieved task prediction for Task ID: 2381 6/162
Successfully retrieved task prediction for Task ID: 2382 7/162
Successfully retrieved task prediction for Task ID: 2360 8/162
Successfully retrieved task prediction for Task ID: 2383 9/162
Successfully retrieved task prediction for Task ID: 2384 10/162
Successfully retrieved task prediction for Task ID: 2359 11/162
Successfully retrieved task prediction for Task ID: 2385 12/162
Successfully retrieved task prediction for Task ID: 2386 13/162
Successfully retrieved task prediction for Task ID: 2387 14/162
Successfully retrieved task prediction for Task ID: 2358 15/162
Successfully retrieved task prediction for Task ID: 2388 16/162
Successfully retrieved task prediction for Task ID: 2389 17/162
Successfully retrieved task prediction for Task ID: 2390 18/162
Successfully retrieved task prediction for Task ID: 2357 19/162
Successfully retrieved task prediction for Task ID: 2391 20/162
...

 

3. 결과 확인