import streamlit as st import torch import requests import os import json from transformers import AutoImageProcessor, AutoModelForImageClassification, AutoConfig from peft import PeftModel, PeftConfig from PIL import Image import torch.nn.functional as F import io # 1. 페이지 기본 설정 (가장 먼저 호출해야 함) st.set_page_config(page_title="Pokemon Classifier", page_icon="🐾", layout="wide", initial_sidebar_state="collapsed") # 2. 커스텀 CSS 적용 (더 서비스스러운 느낌을 위해) st.markdown("""

Pokemon Classifier

""", unsafe_allow_html=True) # 3. 만능 모델 로드 함수 (캐싱 & 에러 방어 적용) @st.cache_resource(show_spinner=False) def load_model(model_path): try: # 허깅페이스 리포지토리 또는 로컬 경로에서 PEFT 설정 파일 확인 try: config = PeftConfig.from_pretrained(model_path) is_peft = True except: is_peft = False if is_peft: base_model_name = config.base_model_name_or_path try: processor = AutoImageProcessor.from_pretrained(model_path) except: processor = AutoImageProcessor.from_pretrained(base_model_name) # 1. 포켓몬 분류는 무조건 150개의 클래스입니다. num_labels = 150 # 2. 라벨 매핑 딕셔너리 구성 (직접 생성하여 충돌 방지) try: # 로컬에 저장된 Full FT 모델의 config가 있다면 가장 완벽한 포켓몬 이름 딕셔너리 사용 config_path = "./saved_model/best_vit_full/config.json" with open(config_path, "r", encoding="utf-8") as f: full_config = json.load(f) id2label = {int(k): v for k, v in full_config["id2label"].items()} label2id = full_config["label2id"] except Exception: try: # 로컬 파일이 없으면 배포된 허깅페이스 저장소에서 올바른 매핑을 강제로 가져옵니다. reference_config = AutoConfig.from_pretrained("gyann/pokemon-vit-full") id2label = {int(k): v for k, v in reference_config.id2label.items()} label2id = reference_config.label2id except Exception: # 최후의 수단으로 더미 생성 id2label = {i: f"LABEL_{i}" for i in range(num_labels)} label2id = {f"LABEL_{i}": i for i in range(num_labels)} # 3. Base 모델을 로드할 때 반드시 num_labels를 명시해야 classifier 헤드 사이즈가 150으로 초기화됩니다. base_model = AutoModelForImageClassification.from_pretrained( base_model_name, num_labels=num_labels, id2label=id2label, label2id=label2id, ignore_mismatched_sizes=True ) # 4. 150 사이즈로 맞춰진 Base 모델에 LoRA 가중치 결합 model = PeftModel.from_pretrained(base_model, model_path) else: processor = AutoImageProcessor.from_pretrained(model_path) model = AutoModelForImageClassification.from_pretrained(model_path) # 비-PEFT 모델도 config에 이름이 없는 경우(Label_15 등)를 대비해 매핑을 덮어씌웁니다. if getattr(model.config, "id2label", {}).get(0, "") == "LABEL_0" or getattr(model.config, "id2label", {}).get("0", "") == "LABEL_0": try: reference_config = AutoConfig.from_pretrained("gyann/pokemon-vit-full") model.config.id2label = {int(k): v for k, v in reference_config.id2label.items()} model.config.label2id = reference_config.label2id except Exception: pass device = "cuda" if torch.cuda.is_available() else "cpu" model.to(device) model.eval() return processor, model, device, None except Exception as e: return None, None, None, str(e) # 4. PokeAPI 데이터 호출 함수 (이미지, 타입, 키, 몸무게) @st.cache_data(show_spinner=False) def get_pokemon_data(pokemon_name): try: name_lower = pokemon_name.lower().replace(" ", "-").replace(".", "").replace("'", "") url = f"https://pokeapi.co/api/v2/pokemon/{name_lower}" response = requests.get(url) if response.status_code == 200: data = response.json() return { "artwork_url": data["sprites"]["other"]["official-artwork"]["front_default"], "types": [t["type"]["name"] for t in data["types"]], "height": data["height"] / 10.0, # meters "weight": data["weight"] / 10.0 # kg } return None except: return None MODEL_PATHS = { "ViT Full Fine-tuning": "gyann/pokemon-vit-full", "ViT + LoRA": "gyann/pokemon-vit-lora", "ViT + QLoRA (4-bit)": "gyann/pokemon-vit-qlora", "ResNet50": "gyann/pokemon-resnet50", "ConvNeXt": "gyann/pokemon-convnext", "Swin Transformer": "gyann/pokemon-swin" } # 5. 설정 및 입력 영역 (깔끔한 컨테이너 UI) with st.container(border=True): st.markdown("### ⚙️ 분석 설정") setting_col1, setting_col2 = st.columns([1, 2]) with setting_col1: compare_models = st.toggle("모델 비교 모드 활성화", value=False) st.caption(f"🚀 현재 가속 장치: **{'GPU (CUDA)' if torch.cuda.is_available() else 'CPU'}**") with setting_col2: if compare_models: col_a, col_b = st.columns(2) with col_a: model_name_a = st.selectbox("Model A", list(MODEL_PATHS.keys()), index=0) with col_b: model_name_b = st.selectbox("Model B", list(MODEL_PATHS.keys()), index=1) with st.spinner("모델 A 로딩 중..."): processor_a, model_a, device_a, err_a = load_model(MODEL_PATHS[model_name_a]) with st.spinner("모델 B 로딩 중..."): processor_b, model_b, device_b, err_b = load_model(MODEL_PATHS[model_name_b]) if err_a: st.error(f"Model A 로드 실패: {err_a}") if err_b: st.error(f"Model B 로드 실패: {err_b}") active_models = [ ("Model A", model_name_a, processor_a, model_a, device_a), ("Model B", model_name_b, processor_b, model_b, device_b) ] else: model_name = st.selectbox("분석에 사용할 모델", list(MODEL_PATHS.keys())) with st.spinner("모델 로딩 중..."): processor, model, device, err = load_model(MODEL_PATHS[model_name]) if err: st.error(f"모델 로드 실패: {err}") active_models = [("Result", model_name, processor, model, device)] with st.container(border=True): st.markdown("### 🖼️ 이미지 선택") EXAMPLE_IMAGES = { "직접 파일 업로드하기": None, "예시: 피카츄 (Pikachu)": "https://raw.githubusercontent.com/PokeAPI/sprites/master/sprites/pokemon/other/official-artwork/25.png", "예시: 이상해씨 (Bulbasaur)": "https://raw.githubusercontent.com/PokeAPI/sprites/master/sprites/pokemon/other/official-artwork/1.png", "예시: 꼬부기 (Squirtle)": "https://raw.githubusercontent.com/PokeAPI/sprites/master/sprites/pokemon/other/official-artwork/7.png", "예시: 파이리 (Charmander)": "https://raw.githubusercontent.com/PokeAPI/sprites/master/sprites/pokemon/other/official-artwork/4.png" } # 라디오 버튼을 가로로 배치하여 모던하게 selected_example = st.radio("테스트 방식을 선택하세요", list(EXAMPLE_IMAGES.keys()), horizontal=True, label_visibility="collapsed") image_to_process = None if selected_example == "직접 파일 업로드하기": uploaded_file = st.file_uploader("포켓몬 이미지를 드래그 앤 드롭하세요", type=["jpg", "jpeg", "png"], label_visibility="collapsed") if uploaded_file is not None: image_to_process = Image.open(uploaded_file).convert("RGB") else: example_url = EXAMPLE_IMAGES[selected_example] try: response = requests.get(example_url) if response.status_code == 200: image_to_process = Image.open(io.BytesIO(response.content)).convert("RGB") except Exception as e: st.error(f"예시 이미지를 불러오는 중 오류가 발생했습니다: {e}") # 6. 추론 및 결과 출력 컴포넌트 def predict_and_display(name_prefix, model_name, processor, model, device, image): if model is None: st.warning(f"모델({model_name})이 정상적으로 로드되지 않았습니다.") return inputs = processor(images=image, return_tensors="pt").to(device) with torch.no_grad(): outputs = model(**inputs) logits = outputs.logits probs = F.softmax(logits, dim=-1)[0] top5_prob, top5_catid = torch.topk(probs, 5) cat_id_0 = top5_catid[0].item() top1_label = model.config.id2label.get(cat_id_0, model.config.id2label.get(str(cat_id_0), "Unknown")) top1_score = top5_prob[0].item() poke_data = get_pokemon_data(top1_label) st.markdown(f"
{name_prefix}: {model_name}
", unsafe_allow_html=True) # 메트릭 st.metric(label="🏆 Top-1 Prediction", value=top1_label.title(), delta=f"{top1_score*100:.1f}% Confidence", delta_color="normal") tab1, tab2, tab3 = st.tabs(["📝 Overview", "📊 Stats", "🔍 Top-5 Details"]) with tab1: if poke_data and poke_data["artwork_url"]: st.image(poke_data["artwork_url"], use_container_width=True) else: st.info("공식 일러스트가 없습니다.") with tab2: if poke_data: st.markdown(f"**속성(Types)**: {', '.join(poke_data['types']).title()}") st.markdown(f"**신장(Height)**: {poke_data['height']} m") st.markdown(f"**체중(Weight)**: {poke_data['weight']} kg") else: st.warning("도감 정보를 가져올 수 없습니다.") with tab3: for i in range(top5_prob.size(0)): c_id = top5_catid[i].item() lbl = model.config.id2label.get(c_id, model.config.id2label.get(str(c_id), "Unknown")) scr = top5_prob[i].item() st.caption(f"**{i+1}. {lbl.title()}**") st.progress(scr, text=f"{scr * 100:.1f}%") # 7. 메인 실행 부 if image_to_process is not None: st.markdown("---") with st.container(): if compare_models: col_img, col_res = st.columns([1, 2.5]) with col_img: st.markdown("### 📷 입력 이미지") st.image(image_to_process, use_container_width=True) with col_res: st.markdown("### ✨ 분석 결과") col1, col2 = st.columns(2) with col1: with st.container(border=True): m_prefix, m_name, m_proc, m_model, m_dev = active_models[0] with st.spinner('분석 중...'): predict_and_display(m_prefix, m_name, m_proc, m_model, m_dev, image_to_process) with col2: with st.container(border=True): m_prefix, m_name, m_proc, m_model, m_dev = active_models[1] with st.spinner('분석 중...'): predict_and_display(m_prefix, m_name, m_proc, m_model, m_dev, image_to_process) else: col1, col2 = st.columns([1, 1.5]) with col1: st.markdown("### 📷 입력 이미지") st.image(image_to_process, use_container_width=True) with col2: st.markdown("### ✨ 분석 결과") with st.container(border=True): m_prefix, m_name, m_proc, m_model, m_dev = active_models[0] with st.spinner('분석 중...'): predict_and_display(m_prefix, m_name, m_proc, m_model, m_dev, image_to_process) # 8. 모델 아키텍처 정보 UI (최하단 배치) ARCHITECTURE_INFO = { "ViT Full Fine-tuning": { "title": "Vision Transformer (ViT) - Full FT", "desc": "이미지를 16x16 픽셀 패치(Patch)로 나누어 처리합니다. 이미지 전체의 **전역적 문맥(Global Context)** 을 파악하는 데 뛰어나며, 모든 가중치를 재학습하여 최고 성능을 도출하지만 연산 비용이 가장 큽니다." }, "ViT + LoRA": { "title": "ViT with LoRA (Low-Rank Adaptation)", "desc": "거대한 원본 모델을 얼려두고(Freeze), 핵심 연산층에 아주 얇은 **학습 가능한 우회로(저랭크 행렬)** 를 덧붙입니다. 단 1%의 파라미터만 학습하여 비용을 극적으로 낮추면서도 Full FT와 유사한 성능을 냅니다." }, "ViT + QLoRA (4-bit)": { "title": "ViT with QLoRA (Quantized LoRA)", "desc": "LoRA에서 한 발 더 나아가, 원본 모델을 **4-bit 정밀도** 로 압축(Quantization)하여 메모리에 적재합니다. VRAM이 매우 부족한 환경에서도 대규모 모델을 튜닝할 수 있게 하는 최적화 기법입니다." }, "ResNet50": { "title": "ResNet50 (Baseline CNN)", "desc": "합성곱(Convolution) 필터를 겹쳐 이미지의 **국소적 패턴(Local Feature)** 을 찾아내는 전통의 강자입니다. 잔차 연결(Residual Connection)로 깊은 신경망의 학습 안정성을 보장합니다." }, "ConvNeXt": { "title": "ConvNeXt (Modernized CNN)", "desc": "트랜스포머의 설계 철학(큰 커널, LayerNorm, GELU 등)을 역으로 CNN에 도입한 **'모던 CNN'** 입니다. CNN의 지역적 귀납적 편향을 유지하면서도 트랜스포머급 성능을 냅니다." }, "Swin Transformer": { "title": "Swin Transformer (Hierarchical ViT)", "desc": "CNN처럼 작은 영역(Window)부터 점점 넓은 영역으로 **계층적(Hierarchical)** 으로 병합하며 학습하는 트랜스포머입니다. ViT가 놓치기 쉬운 미세한 디테일 구분에 강합니다." } } st.markdown("---") st.markdown("#### 🧠 선택된 아키텍처 알아보기") if compare_models: col_info1, col_info2 = st.columns(2) with col_info1: with st.expander(f"{model_name_a} 구조", expanded=False): info = ARCHITECTURE_INFO.get(model_name_a, {"title": model_name_a, "desc": "설명이 준비되지 않았습니다."}) st.markdown(f"**{info['title']}**\n\n{info['desc']}") with col_info2: with st.expander(f"{model_name_b} 구조", expanded=False): info = ARCHITECTURE_INFO.get(model_name_b, {"title": model_name_b, "desc": "설명이 준비되지 않았습니다."}) st.markdown(f"**{info['title']}**\n\n{info['desc']}") else: with st.expander(f"{model_name} 구조", expanded=False): info = ARCHITECTURE_INFO.get(model_name, {"title": model_name, "desc": "설명이 준비되지 않았습니다."}) st.markdown(f"**{info['title']}**\n\n{info['desc']}")