프로젝트/ML 웹앱 배포 전주기

ML 서비스 웹 배포 1 - 로컬, React & FastAPI

ryanhearts 2024. 1. 24. 00:11

결과물

샘플

 


개요

앞서 커스텀 데이터로 반려동물 식별 후 segmentation 작업을 진행했다.

 

커스텀 데이터로 Yolo V8 Fine-Tuning

개요 ML 서비스 웹앱 배포 전주기 체험 프로젝트의 일환 데이터셋 수집, 정제, 모델학습 등 여러 작업을 번거롭게 하지 않아도 쉽게 가져와서 쓸수있는 뛰어난 모델들(메타의 SAM)이 많지만 전주

ryanheart.tistory.com

최종 목표인 반려동물 색상 추출 서비스는 배경의 영향을 제거하기 위해 앞서 얻은 학습 모델을 통해 segmentation 진행 후 진행된다.

이후 완성된 기능은 FastAPI로 구현한 API를 통해 엔드포인트에서 사용하게된다.

웹쪽은 잘 몰라서 ChatGPT한테 물어보면서 React로 간단히 꾸렸다.

전체적인 흐름도

1. React에서 이미지 파일 업로드

2. API에서 수신한 이미지 버퍼 수신

3. segmentation (yolo v8) - 마스킹 이미지 생성

4. color extract (python extcolors library) - 컬러 정보

5. 마스킹된 이미지 인코딩 - numpy array > binary jpeg > base64

6. 인코딩된 이미지 + 컬러정보 json response 전송


API 

ColorExtractor는 내가 만든 클래스로 흐름도상 3, 4번 작업을 진행

해당 클래스의 핵심 메서드

predict : 입력된 이미지 yolo model predict 진행 후 결과 객체 반환

getMask : 이미지, 결과 객체 입력받아 마스크 정보(binary array) 반환

extractFromImageMaskPair : 이미지, 마스크 정보를 받아 마스크 제외한 영역에서 색상 추출

app = FastAPI()
cer = ColorExtractor()

@app.post("/upload")
async def create_upload_file(file: UploadFile = File(...)):
    file_path = f"{MEDIA_ROOT}/{file.filename}"

    img = np.frombuffer(file.file.read(), dtype=np.uint8)
    img = cv2.imdecode(img, cv2.IMREAD_COLOR)

    fp_png = f'{file_path.split(".")[0]}.png'
    cv2.imwrite(fp_png, img)

    pred = cer.predict(img)
    mask = cer.getMask(img, pred)

    if mask is None:
        return {"filename": file.filename, "file_path": file_path, "colors": {}}

    cmap = cer.extractFromImageMaskPair(img, mask)

    thumbnail = cer.getThumbnail(img, mask, pred)
    _, thumbnail_buffer = cv2.imencode('.jpg', thumbnail)
    thumbnail_bytes = base64.b64encode(thumbnail_buffer)

    return {"filename": file.filename, "file_path": file_path, "colors": cmap, "imageBase64": thumbnail_bytes}

React는 코드는 원하는 대략적인 레이아웃과 API 입출력 정보 등을 프롬프트에 담아 ChatGPT 초벌을 짜달라고 했다.

초벌을 받고 2-3번정도 수정하니 잘 동작. 웹 이해도가 높지 않아 리뷰는 생략.

 

이후 진행할 2개 과정

1. 외부에서 접근하기 위한 포트포워딩

2. AWS S3 & EC2 인스턴스를 사용해 배포