Grounding-DINO FastAPI 구현

Machine Learning

Grounding-DINO FastAPI 구현

blog507 2024. 12. 6. 09:46

오늘은 Grounding-DINO 코드를 활용하여 FastAPI를 구현했습니다.

이전 글에도 적어놓았지만 Grounding-DINO는 multi-modal 모델입니다. 이미지와 텍스트를 입력하여 원하는 결과를 얻을 수 있게 됩니다. 이를 잘 활용한다면 귀찮은 annotation 수작업 없이 모델 하나만을 통해 모든 데이터를 1차적으로 annotation할 수 있는 편리함을 가지게 됩니다.

그런데 매번 쓸 때마다 다시 다운로드 받고 설정하는 게 귀찮아서 API로 만들면 어떨까 하는 생각에 만들게 되었습니다.

한번 만들어놓고 잘 활용한다면 정말 좋은 API고 자동으로 labeling도 해줄 수 있어서 편리합니다.

세부적인 코드 주석은 달지 않았지만 틀린 부분이 있으면 언제든 댓글로 말씀해주시면 수정하겠습니다.

논문에 대한 글은 블로그 다른 글에 작성되어있습니다.

2024.07.12 - [Machine Learning/Detection] - Grounding DINO: Marrying DINO with Grounded Pre-Training for Open-Set Object Detection

1. Git Clone 하기 & 필요한 패키지 설치하기

어려운 부분이 아니니 추가 설명은 하지 않겠습니다. Conda나 venv와 같은 가상환경에서 하시는 것은 자유입니다.

git clone https://github.com/IDEA-Research/GroundingDINO.git
cd GroundingDINO

! pip install -r requirements.py
! pip install fastapi uvicorn python-multipart

2. Utils.py 생성

API를 만드는 데 필요한 코드를 생성합니다. utils라는 폴더를 생성 후 utils.py 파일 안에 밑의 코드를 복사 붙여넣기 하시면 됩니다. 함수는 크게 3개가 있고 코드 가독성 및 분리성을 위해서라면 get_transform을 다른 폴더에 두는 게 좋지만 일단 그대로 두겠습니다.

export_json : 이미지를 추론하고 json 파일로 저장 (추후에 원하시는 데이터 형태로 변경하셔도 됩니다.)
visualize_dino : 이미지를 추론하고 나온 결과(bbox)를 이미지에 그림
get_transform : 이미지 transform

! mkdir utils

# utils.py
def export_json(boxes, logits, phrases, img_name, img_width, img_height):
    def convert_to_absolute(box, img_width, img_height):
        cx, cy, w, h = box
        cx = int(cx * img_width)
        cy = int(cy * img_height)
        w = int(w * img_width)
        h = int(h * img_height)
        xmin = cx - w // 2
        ymin = cy - h // 2
        xmax = cx + w // 2
        ymax = cy + h // 2
        xmin = max(0, min(xmin, img_width - 1))
        ymin = max(0, min(ymin, img_height - 1))
        xmax = max(0, min(xmax, img_width - 1))
        ymax = max(0, min(ymax, img_height - 1))
        return xmin, ymin, xmax, ymax

    output_list = []

    boxes = boxes.tolist() if isinstance(boxes, torch.Tensor) else boxes
    logits = logits.tolist() if isinstance(logits, torch.Tensor) else logits

    for box, logit, phrase in zip(boxes, logits, phrases):
        xmin, ymin, xmax, ymax = convert_to_absolute(box, img_width, img_height)
        box_info = {
            "label": phrase,
            "bbox": {"x1": xmin, "y1": ymin, "x2": xmax, "y2": ymax},
            "score": float(logit),
        }
        output_list.append(box_info)

    return output_list


def visualize_dino(image_source, boxes, logits, phrases) -> None:
    image_visualized = image_source.copy()
    height, width, _ = image_source.shape

    for box, logit, phrase in zip(boxes, logits, phrases):
        cx, cy = int(box[0] * width), int(box[1] * height)
        box_width, box_height = int(box[2] * width), int(box[3] * height)
        x1, y1 = int(cx - box_width / 2), int(cy - box_height / 2)
        x2, y2 = int(cx + box_width / 2), int(cy + box_height / 2)

        cv2.rectangle(image_visualized, (x1, y1), (x2, y2), (0, 255, 0), 2)

        label = f"{phrase} ({logit:.2f})"
        font = cv2.FONT_HERSHEY_SIMPLEX
        font_scale = 0.5
        font_thickness = 1
        text_size, _ = cv2.getTextSize(label, font, font_scale, font_thickness)
        text_x, text_y = x1, y1 - 10 if y1 - 10 > 10 else y1 + 10 + text_size[1]

        cv2.rectangle(
            image_visualized, (text_x, text_y - text_size[1] - 2), (text_x + text_size[0], text_y + 2), (0, 255, 0), -1
        )
        cv2.putText(image_visualized, label, (text_x, text_y), font, font_scale, (255, 255, 255), font_thickness)

    return image_visualized


def get_transform(image):
    image = Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))

    transform = T.Compose(
        [
            T.RandomResize([800], max_size=1333),
            T.ToTensor(),
            T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
        ]
    )
    image_transformed, _ = transform(image, None)

    return image_transformed

3. app.py 생성

마지막 코드 추가입니다. API는 크게 json 파일을 출력하는 것과 이미지를 출력하는 것 두개로 나뉘어져있습니다. 모델을 로드한 후에 이미지와 텍스트를 입력하면 결과가 나오게 되는 로직입니다. 코드가 어렵지 않으니 이해하는 것도 어렵지 않으실 껍니다.

# app.py
import io
import numpy as np
import cv2
from fastapi import FastAPI, Form, UploadFile, File
from fastapi.responses import StreamingResponse, JSONResponse
from groundingdino.util.inference import load_model, predict
import os
from typing import *
import sys

sys.path.append(os.path.join(os.path.dirname(__file__), "./utils"))
from utils import export_json, visualize_dino, get_transform

app = FastAPI()

# 모델 로드 경로 설정
HOME = "./GroundingDINO"
CONFIG_PATH = os.path.join(HOME, "groundingdino/config/GroundingDINO_SwinT_OGC.py")
WEIGHTS_PATH = os.path.join(HOME, "weights", "groundingdino_swint_ogc.pth")


# 애플리케이션 시작 시 모델 로드
@app.on_event("startup")
def startup_event():
    global model
    model = load_model(CONFIG_PATH, WEIGHTS_PATH)


@app.post("/grounding-dino")
async def grounding_dino(
    text_prompt: str = Form(..., example="Detect objects"),
    box_threshold: float = Form(0.35, example=0.35),
    text_threshold: float = Form(0.25, example=0.25),
    image_file: UploadFile = File(..., example="image.jpg"),
):

    image_data = await image_file.read()
    image_array = np.frombuffer(image_data, np.uint8)
    image_source = cv2.imdecode(image_array, cv2.IMREAD_COLOR)

    image_transformed = get_transform(image_source)

    boxes, logits, phrases = predict(
        model=model,
        image=image_transformed,
        caption=text_prompt,
        box_threshold=box_threshold,
        text_threshold=text_threshold,
    )

    img_height, img_width = image_source.shape[:2]
    json_output = export_json(boxes, logits, phrases, image_file.filename, img_width, img_height)

    return JSONResponse(content=json_output)


@app.post("/grounding-dino/visualize")
async def visualize_grounding_dino(
    text_prompt: str = Form(..., example="Detect objects"),
    box_threshold: float = Form(0.35, example=0.35),
    text_threshold: float = Form(0.25, example=0.25),
    image_file: UploadFile = File(..., example="image.jpg"),
):
    image_data = await image_file.read()
    image_array = np.frombuffer(image_data, np.uint8)
    image_source = cv2.imdecode(image_array, cv2.IMREAD_COLOR)

    image_transformed = get_transform(image_source)

    boxes, logits, phrases = predict(
        model=model,
        image=image_transformed,
        caption=text_prompt,
        box_threshold=box_threshold,
        text_threshold=text_threshold,
    )
    image_visualized = visualize_dino(image_source, boxes, logits, phrases)

    _, img_encoded = cv2.imencode(".png", image_visualized)
    img_bytes = io.BytesIO(img_encoded.tobytes())

    return StreamingResponse(img_bytes, media_type="image/png")

4. 실행 및 결과 확인

실행하는 방법은 아래 커맨드를 입력해주시면 됩니다. 그러면 웹에서 localhost:4444에 접근을 할 수 있게 됩니다. 근데 이 화면은 아무 UI가 없으니 확인하려면 Swagger로 확인해보시는 게 좋습니다.

! uvicorn app:app --host 0.0.0.0 --port 4444 # port는 자유

그래서 localhost:4444/docs로 접근해보죠. 그러면 아래와 같은 화면이 뜨게 됩니다.

어떻게 사용하는 지 알려드리겠습니다.

일단 두 API 중 위부터 보겠습니다. 화살표를 누르고 "Try it out" 버튼을 누르면 아래와 같은 화면이 나오는데 여기서 text_prompt와 image_file을 업로드 하시고 Execute를 누르면 됩니다. (box_threshold와 text_threshold는 수정하셔도 되고 안하셔도 됩니다.)

실행하면 위와 같이 정보를 얻을 수 있게 됩니다. 즉, annotation을 자동으로 해주는 거죠. 물론 .json파일로 다운로드도 할 수 있습니다! (우측 하단 참고)

아래 API도 똑같이 실행해보겠습니다.

꽤 간단하고 빠르게 이미지를 출력하는 것을 보실 수 있습니다.

저는 이미지 한장만을 이용해서 진행했지만 이미지 폴더를 받아서 폴더 내 모든 이미지를 이렇게 annotations 할 수 있다면 훨씬 더 유용하게 쓸 수 있을 것이라 생각합니다.

궁금하신 사항 있으시면 언제든 댓글로 의견 주시면 감사하겠습니다.

'Machine Learning' 카테고리의 다른 글

PDFTranslate FastAPI 구현 (1)	2025.02.04
Stable Diffusion FastAPI 구현 (0)	2025.02.03
GLIP : Grounded Language-Image Pre-training (1)	2024.07.26
CNN 기반 모델들 (0)	2024.06.25
딥러닝 기초 지식 (2)	2024.06.11

현재글Grounding-DINO FastAPI 구현

ai-notes

Kim, Lim,

일	월	화	수	목	금	토
			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

ai-notes