[AI] Flask로 YOLOv5 모델 서비스

Bora Kwon·2022년 7월 26일
0

오늘 바짝 집중해서 웹을 어느 정도 만들었다.
OCR만 연동하면 이번 주에 하기로 한 건 거의 끝나는 셈

어제 돌린 Yolov7은 아직 돌아가고 있어서...
Yolo7은 연동할 수 있을지 모르겠지만 일단
Azure ML Studio에서도 작업했고 (Yolov7 돌리고 있어서 그런지 SQlite가 계속 잠김)
로컬에서도 작업했고 (그래서 그냥 로컬에서 작업 후)
Goorm IDE에서도 작업해서 복잡하지만 (Goorm IDE로 업로드함)

SQLite3 활용한 Database 연동

Goorm IDE에는 기본적으로 SQLite 설치되어 있지 않음 (Flask 선택해서 띄웠는데도)

SQLite3 설치

apt-get update
apt-get install sqlite3

DB 만들어주기

$ sqlite3
sqlite> .open users.sqlite3
sqlite> .exit
$ ls
users.sqlite3
$ sqlite3 users.sqlite3
sqlite> 

테이블 생성

CREATE TABLE users (
id INTEGER PRIMARY KEY AUTOINCREMENT NOT NULL,
name TEXT NOT NULL,
carNumber TEXT NOT NULL
);

간단하게 id, 이름, 차량 번호만 저장하도록 작성했다.
차량번호 중복일 때 처리도 해주면 좋은데 나중에 하든가 하고

Flask 연동

import sqlite3

conn = sqlite3.connect('users.sqlite3')
cursor = conn.cursor()
cursor.execute('SELECT * FROM users WHERE carNumber=?', (ocr_result, ))
user = cursor.fetchone()
conn.close()

이런 식으로 사용하면 됨.

Flask로 Yolov5 service

    pt_file_path = "<pt 파일 경로>"
    model = torch.hub.load('ultralytics/yolov5', 'custom', pt_file_path)

torch.hub.load에서 저 yolov5를 가지고 와야 model을 사용할 수 있는 듯..?

    img_bytes = file.read()
    img = Image.open(io.BytesIO(img_bytes))
    results = model(img, size=640)
    results.render()
    
    for img in results.imgs:
        img_base64 = Image.fromarray(img)
        img_base64.save(img_src)

그리고 model(이미지) 주고 render해서 저장하면 박스 친 결과가 나옴.

내일은 OCR 연동해야지

profile
Software Developer

0개의 댓글