jhynsoo/worlabel
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

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

Browse Source 
Fix: classification process_prediction_result() 에러 수정

See merge request s11-s-project/S11P21S002!227
This commit is contained in:
김진현 2024-09-27 16:03:25 +09:00
commit 6bf4bba555
1 changed files with 15 additions and 16 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

31
ai/app/api/yolo/classfication.py
View File

@ -42,39 +42,38 @@ def get_model(project_id:int, model_key:str):
# 추론 결과 처리 함수 # 추론 결과 처리 함수
def process_prediction_result(result, image, label_map): def process_prediction_result(result, image, label_map):
try: try:
label_name = None shapes = []
# top 5에 해당하는 class id 순회 # top 5에 해당하는 class id 순회
for class_id in result.probs.top5: for class_id in result.probs.top5:
name = result.names[class_id] # class id에 해당하는 label_name label_name = result.names[class_id] # class id에 해당하는 label_name
if name in label_map: # name이 사용자 레이블 카테고리에 있을 경우 if label_name in label_map: # name이 사용자 레이블 카테고리에 있을 경우
label_name = name # label_name 설정 shapes = [
Shape(
label=label_name,
color=get_random_color(),
points=[[0.0, 0.0]],
group_id=label_map[label_name],
shape_type='point',
flags={}
)
] # label_name 설정
break break
label_data = LabelData( label_data = LabelData(
version="0.0.0", version="0.0.0",
task_type="cls", task_type="cls",
shapes=[], shapes=shapes,
split="none", split="none",
imageHeight=result.orig_img.shape[0], imageHeight=result.orig_img.shape[0],
imageWidth=result.orig_img.shape[1], imageWidth=result.orig_img.shape[1],
imageDepth=result.orig_img.shape[2] imageDepth=result.orig_img.shape[2]
) )
if label_name: # label_name을 설정한게 있다면 추가
shape = Shape(
label= label_name,
color= get_random_color(),
points= [[0.0, 0.0]],
group_id= label_map[label_name],
shape_type= 'point',
flags= {}
)
LabelData.shapes.append(shape)
return PredictResponse( return PredictResponse(
image_id=image.image_id, image_id=image.image_id,
data=label_data.model_dump_json() data=label_data.model_dump_json()
) )
except KeyError as e: except KeyError as e:
raise HTTPException(status_code=500, detail="KeyError: " + str(e)) raise HTTPException(status_code=500, detail="KeyError: " + str(e))
except Exception as e: except Exception as e: