Merge branch 'ai/feat/image' into 'ai/develop'

Refactor: API endPoint 변경 및 이미지 메모리 처리

See merge request s11-s-project/S11P21S002!43
This commit is contained in:
김진현 2024-09-04 17:44:39 +09:00
commit 635c2a2033
3 changed files with 17 additions and 5 deletions

View File

@ -5,8 +5,7 @@ from services.ai_service import load_detection_model
from typing import List
router = APIRouter()
@router.post("/predict", response_model=List[PredictResponse])
@router.post("/detection", response_model=List[PredictResponse])
def predict(request: PredictRequest):
version = "0.1.0"
@ -20,12 +19,23 @@ def predict(request: PredictRequest):
results = []
try:
for image in request.image_list:
# URL에서 이미지를 메모리로 로드 TODO: 추후 메모리에 할지 어떻게 해야할지 or 병렬 처리 고민
# response = requests.get(image.image_url)
# 이미지 데이터를 메모리로 로드
# img = Image.open(io.BytesIO(response.content))
predict_results = model.predict(
source=image.image_url,
source=image.image_url,
iou=request.iou_threshold,
conf=request.conf_threshold,
classes=request.classes)
classes=request.classes
)
results.append(predict_results[0])
# 메모리에서 이미지 객체 해제
# img.close()
# del img
except Exception as e:
raise HTTPException(status_code=500, detail="model predict exception: "+str(e))
@ -57,6 +67,7 @@ def predict(request: PredictRequest):
}
response.append({
"image_id":image.image_id,
"image_url":image.image_url,
"data":label_data
})
except Exception as e:

View File

@ -4,7 +4,7 @@ from api.yolo.detection import router as yolo_detection_router
app = FastAPI()
# 각 기능별 라우터를 애플리케이션에 등록
app.include_router(yolo_detection_router, prefix="/api/yolo/detection")
app.include_router(yolo_detection_router, prefix="/api")
# 애플리케이션 실행
if __name__ == "__main__":

View File

@ -20,4 +20,5 @@ class LabelData(BaseModel):
class PredictResponse(BaseModel):
image_id: int
image_url: str
data: LabelData