README.md

metadata

title: Three-View-Style-Embedder
emoji: 🎨
colorFrom: blue
colorTo: purple
sdk: gradio
sdk_version: 4.0.0
app_file: app.py
pinned: false

Three-View-Style-Embedder

17,000+ 작가의 일러스트 스타일을 임베딩하고 분류하는 Multi-branch Metric Learning 모델

개요

Three-View-Style-Embedder는 일러스트의 전체 이미지, 얼굴, 눈 세 가지 뷰를 결합하여 작가 고유의 스타일을 512차원 벡터로 임베딩합니다.

아키텍처

┌─────────────────────────────────────────────────────────────┐
│                      Input Images                            │
├───────────────────┬───────────────────┬─────────────────────┤
│    Full Image     │    Face Crop      │      Eye Crop       │
│      (필수)        │     (선택)         │       (선택)         │
└─────────┬─────────┴─────────┬─────────┴──────────┬──────────┘
          │                   │                    │
          ▼                   ▼                    ▼
┌─────────────────────────────────────────────────────────────┐
│            3× EVA02-Large Encoders (각 304M params)          │
└─────────────────────────────────────────────────────────────┘
          │                   │                    │
          └─────────┬─────────┴─────────┬──────────┘
                    │                   │
                    ▼                   ▼
          ┌─────────────────────────────────────┐
          │         Gated Fusion Module          │
          │   (각 브랜치 중요도 동적 학습)         │
          └─────────────────┬───────────────────┘
                            │
                            ▼
          ┌─────────────────────────────────────┐
          │     Style Embedding Head (512-d)     │
          │         L2 Normalized Output         │
          └─────────────────┬───────────────────┘
                            │
               ┌────────────┴────────────┐
               ▼                         ▼
        ┌────────────┐            ┌────────────┐
        │  ArcFace   │            │  Multi-    │
        │   Loss     │            │ Similarity │
        └────────────┘            └────────────┘

프로젝트 구조

Three-View-Style-Embedder/
├── config.py              # 설정 관리
├── model.py               # 모델 아키텍처
├── dataset.py             # 데이터셋 및 데이터로더
├── losses.py              # Loss 함수 (ArcFace, Multi-Similarity, Center)
├── trainer.py             # 학습 로직
├── train.py               # 학습 실행 스크립트
├── evaluate.py            # 평가 스크립트
├── extract_embeddings.py  # 임베딩 추출 스크립트
├── app.py                 # Hugging Face Spaces용 Web UI (자동 다운로드)
├── local_app.py           # 로컬 실행용 Web UI
├── inference_utils.py     # 추론 유틸리티 (공통)
├── requirements.txt       # 의존성
└── README.md              # 문서

설치

pip install -r requirements.txt

Windows에서는 작업 폴더의 .venv를 사용하도록 권장합니다 (PATH에 있는 다른 Python으로 실행하면 Gradio/torch 버전 충돌이 날 수 있습니다).

데이터셋 구조

./dataset/              # 전체 이미지 (필수)
    ├── artist_name_1/
    │   ├── image1.jpg
    │   └── image2.png
    └── artist_name_2/
        └── ...

./dataset_face/         # 얼굴 크롭 (선택)
    └── (동일 구조)

./dataset_eyes/         # 눈 크롭 (선택)
    └── (동일 구조)

사용법

1. 학습

python train.py \
    --dataset_root ./dataset \
    --dataset_face_root ./dataset_face \
    --dataset_eyes_root ./dataset_eyes \
    --epochs 100 \
    --batch_size 256

2. 임베딩 추출

학습된 모델로 모든 작가의 임베딩을 추출:

python extract_embeddings.py \
    --checkpoint ./checkpoints/best_model.pt \
    --output ./embeddings/artist_embeddings.npz \
    --max_combinations 10 \
    --batch_size 256

3. Web UI 실행

로컬 실행

python local_app.py \
    --checkpoint ./checkpoints/best_model.pt \
    --embeddings ./embeddings/artist_embeddings.npz

또는 Windows에서:

run.bat

브라우저에서 http://localhost:7860 접속

Hugging Face Spaces 실행

app.py는 Hugging Face Spaces에서 자동으로 모델을 다운로드하여 실행됩니다. Spaces에 업로드하면 자동으로 작동합니다.

모델 스펙

항목	값
Backbone	EVA02-Large-14-CLIP × 3
Total Parameters	~920M
Embedding Dimension	512
Input Size	224 × 224
Loss	ArcFace + Multi-Similarity + Center

성능

메트릭	값
Top-1 Accuracy	~77%
Top-5 Accuracy	~92%
클래스 수	17,000+

출력 파일

체크포인트 (./checkpoints/)

• best_model.pt: 최고 성능 모델
• checkpoint_epoch_N.pt: 에포크별 체크포인트

임베딩 (./embeddings/)

• artist_embeddings.npz: 작가별 평균 임베딩
• artist_embeddings.json: 메타데이터

API 사용 예시

from extract_embeddings import load_embeddings, find_similar_artists
from app import get_image_embedding

# 임베딩 로드
artist_names, embeddings = load_embeddings('./embeddings/artist_embeddings.npz')

# 이미지에서 임베딩 추출
query_emb = get_image_embedding(model, image, device)

# 유사 작가 검색
similar = find_similar_artists(query_emb, artist_names, embeddings, top_k=10)
for name, score in similar:
    print(f"{name}: {score:.4f}")

라이선스

MIT License