app.py

import gradio as gr
from fastapi import FastAPI, File, UploadFile
from fastapi.responses import JSONResponse
from PIL import Image
import numpy as np
import io

import torch
import torch.nn.functional as F
from transformers import AutoImageProcessor, AutoModelForImageClassification


device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
image_processor = AutoImageProcessor.from_pretrained("wesleyacheng/dog-breeds-multiclass-image-classification-with-vit")
model = AutoModelForImageClassification.from_pretrained("wesleyacheng/dog-breeds-multiclass-image-classification-with-vit").to(device)
labels = {int(k): v for k, v in model.config.id2label.items()}
print(labels)

# FastAPI 앱 생성
app = FastAPI()

# 품종 한글 번역
BREED_TRANSLATIONS = {
    "otterhound": "오터하운드", "cocker_spaniel": "코커 스패니얼", "brittany_spaniel": "브리타니 스패니얼",
    "afghan_hound": "아프간 하운드", "maltese_dog": "몰티즈", "schipperke": "스키퍼키",
    "irish_setter": "아이리시 세터", "pekinese": "페키니즈", "golden_retriever": "골든 리트리버",
    "vizsla": "비즐라", "welsh_springer_spaniel": "웰시 스프링거 스패니얼",
    "staffordshire_bullterrier": "스태퍼드셔 불테리어", "border_collie": "보더 콜리",
    "irish_terrier": "아이리시 테리어", "eskimo_dog": "에스키모견", "pug": "퍼그",
    "kelpie": "켈피", "yorkshire_terrier": "요크셔 테리어", "tibetan_terrier": "티베탄 테리어",
    "walker_hound": "워커 하운드", "affenpinscher": "아펜핀셔", "cardigan": "카디건 코기",
    "english_springer": "잉글리시 스프링거", "english_foxhound": "잉글리시 폭스하운드",
    "west_highland_white_terrier": "웨스트 하이랜드 화이트 테리어", "lakeland_terrier": "레이클랜드 테리어",
    "rhodesian_ridgeback": "로디지안 리지백", "gordon_setter": "고든 세터", "lhasa": "라사압소",
    "curly": "컬리 코티드 리트리버", "beagle": "비글", "tibetan_mastiff": "티베탄 마스티프",
    "sussex_spaniel": "서식스 스패니얼", "saint_bernard": "세인트 버나드", "toy_terrier": "토이 테리어",
    "standard_poodle": "스탠다드 푸들", "bernese_mountain_dog": "버니즈 마운틴 독", "pomeranian": "포메라니안",
    "ibizan_hound": "이비잔 하운드", "redbone": "레드본 쿤하운드", "toy_poodle": "토이 푸들",
    "basset": "바셋 하운드", "scottish_deerhound": "스코티시 디어하운드", "miniature_pinscher": "미니어처 핀셔",
    "basenji": "바센지", "border_terrier": "보더 테리어", "bedlington_terrier": "베들링턴 테리어",
    "kerry_blue_terrier": "케리 블루 테리어", "weimaraner": "와이마라너", "english_setter": "잉글리시 세터",
    "bluetick": "블루틱 쿤하운드", "boston_bull": "보스턴 테리어", "italian_greyhound": "이탈리안 그레이하운드",
    "dandie_dinmont": "댄디 딘몬트 테리어", "airedale": "에어데일 테리어", "irish_water_spaniel": "아이리시 워터 스패니얼",
    "norfolk_terrier": "노퍽 테리어", "wire": "와이어 폭스 테리어", "french_bulldog": "프렌치 불독",
    "soft": "소프트 코티드 휘튼 테리어", "komondor": "코몬도르", "african_hunting_dog": "아프리카 들개",
    "siberian_husky": "시베리안 허스키", "newfoundland": "뉴펀들랜드", "bouvier_des_flandres": "부비에 데 플랑드르",
    "saluki": "살루키", "shetland_sheepdog": "셰틀랜드 시프독", "collie": "콜리", "rottweiler": "로트와일러",
    "silky_terrier": "실키 테리어", "norwegian_elkhound": "노르웨이 엘크하운드", "chihuahua": "치와와",
    "leonberg": "레온베르거", "norwich_terrier": "노리치 테리어", "cairn": "케언 테리어", "boxer": "복서",
    "borzoi": "보르조이", "dhole": "승냥이", "samoyed": "사모예드", "german_shepherd": "저먼 셰퍼드",
    "labrador_retriever": "래브라도 리트리버", "blenheim_spaniel": "블레넘 스패니얼", "groenendael": "그로넨달",
    "doberman": "도베르만", "great_dane": "그레이트 데인", "flat": "플랫 코티드 리트리버",
    "appenzeller": "아펜젤러", "shih": "시추", "japanese_spaniel": "재패니즈 스패니얼",
    "greater_swiss_mountain_dog": "그레이터 스위스 마운틴 독", "black": "블랙 앤 탄 쿤하운드",
    "dingo": "딩고", "great_pyrenees": "그레이트 피레니즈", "whippet": "휘펫", "keeshond": "키스혼트",
    "malinois": "말리노이즈", "american_staffordshire_terrier": "아메리칸 스태퍼드셔 테리어",
    "mexican_hairless": "멕시칸 헤어리스", "giant_schnauzer": "자이언트 슈나우저",
    "brabancon_griffon": "브라방송 그리펑", "kuvasz": "쿠바스", "miniature_poodle": "미니어처 푸들",
    "irish_wolfhound": "아이리시 울프하운드", "briard": "브리아드", "clumber": "클럼버 스패니얼",
    "standard_schnauzer": "스탠다드 슈나우저", "bull_mastiff": "불 마스티프", "malamute": "말라뮤트",
    "sealyham_terrier": "실리엄 테리어", "entlebucher": "엔틀부처", "chow": "차우차우",
    "papillon": "파피용", "pembroke": "펨브로크 코기", "german_short": "저먼 쇼트헤어드 포인터",
    "old_english_sheepdog": "올드 잉글리시 시프독", "chesapeake_bay_retriever": "체서피크 베이 리트리버",
    "scotch_terrier": "스카치 테리어", "australian_terrier": "오스트레일리안 테리어",
    "miniature_schnauzer": "미니어처 슈나우저", "bloodhound": "블러드하운드"
}


def predict(image: Image.Image, k=10):
    inputs = image_processor(images=image, return_tensors="pt").to(device)

    with torch.no_grad():
        outputs = model(**inputs)

    probs = F.softmax(outputs.logits, dim=-1)[0]

    top_probs, top_idxs = torch.topk(probs, k=k)

    top_probs = top_probs.cpu().tolist()
    top_idxs = top_idxs.cpu().tolist()

    # 영문 품종명을 한글로 번역하여 반환 (확률에 100 곱해서 퍼센트로)
    results = {}
    for i, prob in zip(top_idxs, top_probs):
        breed_name = labels[i]
        korean_name = BREED_TRANSLATIONS.get(breed_name, breed_name)
        results[korean_name] = round(float(prob) * 100, 3)

    return results


def classify(profileImage: Image.Image):
    result = predict(profileImage)
    return result


# FastAPI 엔드포인트: FormData로 이미지 받기
@app.post("/api/classify")
async def classify_api(profileImage: UploadFile = File(...)):
    try:
        # 이미지 파일 읽기
        contents = await profileImage.read()
        image = Image.open(io.BytesIO(contents))

        # 예측 수행
        result = predict(image)

        return JSONResponse(content={"success": True, "data": result})
    except Exception as e:
        return JSONResponse(
            content={"success": False, "error": str(e)},
            status_code=400
        )


# Gradio 인터페이스
demo = gr.Interface(fn=classify, inputs=gr.Image(type="pil", label="profileImage"), outputs="label")

# Gradio를 FastAPI에 마운트
app = gr.mount_gradio_app(app, demo, path="/")

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860)