import io
import torch
from PIL import Image
from fastapi import APIRouter, File, UploadFile
from torchvision.models import resnet50, ResNet50_Weights
router = APIRouter()
print("🍳 모델(ResNet50)을 로딩 중입니다...")
# ResNet의 지식을 가져옵니다 가중치는 지능 DEFAULT는 가장 최신버전을 가져와라
weights = ResNet50_Weights.DEFAULT
model = resnet50(weights=weights)
# 모델을 실전모드로 전환한다
model.eval()
#resnet은 224크기의 이미지밖에 못쓰기때문에 정규화에 맞아야한다 224로 이미지를 맞춰라
preprocess = wei resghts.transforms()
@router.post("/classify")
async def classify_image(file:UploadFile = File(...)):
#0과 1로된 데이터를 읽고
image_data = await file.read()
#메모리에 있는 데이터를 파일로 포장해
image = Image.open(io.BytesIO(image_data)).convert("RGB")
# 이미지를 포장해서 input_tensor에 넣고
input_tensor = preprocess(image)
# ai모델은 항상 여러장을 처리하게 설계되어있어서 (batch)라고함 사진 1장은 3,224,224면 한장은
# 모델이 받지 않기때문 unsqueeze를 통해서 1,3,224,224로 만들어서 1장짜리 묶음으로 포장해서 줍니다
input_batch = input_tensor.unsqueeze(0)
# 파이토치는 계산할때마다 오답노트때문에 기록하는데 메모리를 많이먹어서 no_grad()로 기록하지마라고합니다
with torch.no_grad():
output = model(input_batch)
#softmax는 고양이2점 강아지 10점 이런걸 도합 100%가되게만들어서 강아지 85% 고양이 15% 로 표현하게 만듭니다
probabilities = torch.nn.functional.softmax(output[0], dim=0)
# topk은 top n의 갯수를 뽑아라라는말입니다
top3_prob, top3_catid = torch.topk(probabilities, 3)
results = []
for i in range(top3_prob.size(0)):
results.append({
"class_name" : weights.meta["categories"][top3_catid[i]],
"confidence" : f"{top3_prob[i].item() * 100:2f}%"
})
return {"results": results}