""" Three-View-Style-Embedder - Fast Embedding Extraction DataLoader 기반 고속 배치 처리 """ import argparse import random import itertools from pathlib import Path from typing import Dict, List, Optional, Tuple import json import torch import torch.nn.functional as F from torch.utils.data import Dataset, DataLoader from torchvision import transforms from PIL import Image import numpy as np from tqdm import tqdm from collections import defaultdict from config import get_config from model import ArtistStyleModel class ArtistCombinationDataset(Dataset): """모든 작가의 조합을 미리 생성하는 데이터셋""" def __init__( self, dataset_root: str, dataset_face_root: str, dataset_eyes_root: str, max_combinations: int = 30, ): self.dataset_root = Path(dataset_root) self.dataset_face_root = Path(dataset_face_root) self.dataset_eyes_root = Path(dataset_eyes_root) self.max_combinations = max_combinations self.transform = transforms.Compose([ transforms.Resize((224, 224)), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ]) # 모든 샘플 미리 생성: (artist_name, full_path, face_path, eye_path) self.samples = [] self.artist_to_indices = defaultdict(list) # artist -> [sample indices] self._build_samples() def _get_image_paths(self, folder: Path) -> List[Path]: if not folder.exists(): return [] return list(folder.glob("*.jpg")) + list(folder.glob("*.png")) + list(folder.glob("*.webp")) def _build_samples(self): """모든 작가의 모든 조합 미리 생성""" print("Building sample combinations...") artist_dirs = [d for d in self.dataset_root.iterdir() if d.is_dir()] for artist_dir in tqdm(artist_dirs, desc="Preparing artists"): artist_name = artist_dir.name full_paths = self._get_image_paths(artist_dir) if not full_paths: continue face_paths = self._get_image_paths(self.dataset_face_root / artist_name) eye_paths = self._get_image_paths(self.dataset_eyes_root / artist_name) # 조합 생성 face_options = face_paths if face_paths else [None] eye_options = eye_paths if eye_paths else [None] all_combinations = list(itertools.product(full_paths, face_options, eye_options)) random.shuffle(all_combinations) selected = all_combinations[:self.max_combinations] # 샘플 추가 for full_path, face_path, eye_path in selected: idx = len(self.samples) self.samples.append((artist_name, full_path, face_path, eye_path)) self.artist_to_indices[artist_name].append(idx) print(f"Total samples: {len(self.samples)} from {len(self.artist_to_indices)} artists") def _load_image(self, path: Optional[Path]) -> Optional[torch.Tensor]: if path is None: return None try: img = Image.open(path) if img.mode in ('RGBA', 'LA', 'P'): background = Image.new('RGB', img.size, (255, 255, 255)) if img.mode == 'P': img = img.convert('RGBA') if img.mode in ('RGBA', 'LA'): background.paste(img, mask=img.split()[-1]) img = background else: img = img.convert('RGB') else: img = img.convert('RGB') return self.transform(img) except: return None def __len__(self): return len(self.samples) def __getitem__(self, idx): artist_name, full_path, face_path, eye_path = self.samples[idx] full_tensor = self._load_image(full_path) if full_tensor is None: full_tensor = torch.zeros(3, 224, 224) valid = False else: valid = True face_tensor = self._load_image(face_path) has_face = face_tensor is not None if not has_face: face_tensor = torch.zeros(3, 224, 224) eye_tensor = self._load_image(eye_path) has_eye = eye_tensor is not None if not has_eye: eye_tensor = torch.zeros(3, 224, 224) return { 'full': full_tensor, 'face': face_tensor, 'eye': eye_tensor, 'has_face': has_face, 'has_eye': has_eye, 'artist': artist_name, 'valid': valid, 'idx': idx, } def collate_fn(batch): return { 'full': torch.stack([x['full'] for x in batch]), 'face': torch.stack([x['face'] for x in batch]), 'eye': torch.stack([x['eye'] for x in batch]), 'has_face': torch.tensor([x['has_face'] for x in batch]), 'has_eye': torch.tensor([x['has_eye'] for x in batch]), 'artist': [x['artist'] for x in batch], 'valid': torch.tensor([x['valid'] for x in batch]), 'idx': torch.tensor([x['idx'] for x in batch]), } @torch.no_grad() def extract_all_embeddings( checkpoint_path: str, dataset_root: str, dataset_face_root: str, dataset_eyes_root: str, output_path: str, max_combinations: int = 30, batch_size: int = 64, num_workers: int = 8, device: str = 'cuda', ): """고속 임베딩 추출""" requested_device = device if requested_device.startswith('cuda') and not torch.cuda.is_available(): print( "[WARN] --device=cuda requested but torch.cuda.is_available() is False. " "Falling back to CPU. (Install a CUDA-enabled PyTorch build to use GPU.)" ) requested_device = 'cpu' device = torch.device(requested_device) # 모델 로드 print("Loading model...") # Always load checkpoint on CPU to avoid duplicating large tensors on GPU. checkpoint = torch.load(checkpoint_path, map_location='cpu') config = get_config() model = ArtistStyleModel( num_classes=len(checkpoint['artist_to_idx']), embedding_dim=config.model.embedding_dim, hidden_dim=config.model.hidden_dim, ) model.load_state_dict(checkpoint['model_state_dict']) # Reduce VRAM: keep weights in FP16 on CUDA. if device.type == 'cuda': model = model.to(dtype=torch.float16) model = model.to(device) model.eval() # 데이터셋 생성 dataset = ArtistCombinationDataset( dataset_root=dataset_root, dataset_face_root=dataset_face_root, dataset_eyes_root=dataset_eyes_root, max_combinations=max_combinations, ) dataloader = DataLoader( dataset, batch_size=batch_size, shuffle=False, num_workers=num_workers, collate_fn=collate_fn, pin_memory=True, ) # 임베딩 저장용 all_embeddings = torch.zeros(len(dataset), config.model.embedding_dim) all_valid = torch.zeros(len(dataset), dtype=torch.bool) # 배치 추론 (AMP 사용) print("Extracting embeddings with AMP...") for batch in tqdm(dataloader, desc="Processing"): full = batch['full'].to(device) face = batch['face'].to(device) eye = batch['eye'].to(device) has_face = batch['has_face'].to(device) has_eye = batch['has_eye'].to(device) indices = batch['idx'] valid = batch['valid'] with torch.cuda.amp.autocast(enabled=(device.type == 'cuda')): embeddings = model.get_embeddings(full, face, eye, has_face, has_eye) all_embeddings[indices] = embeddings.float().cpu() all_valid[indices] = valid # 작가별 평균 계산 print("Computing artist averages...") artist_embeddings = {} failed_artists = [] for artist_name, indices in tqdm(dataset.artist_to_indices.items(), desc="Averaging"): indices = torch.tensor(indices) valid_mask = all_valid[indices] if valid_mask.sum() == 0: failed_artists.append(artist_name) continue valid_embeddings = all_embeddings[indices][valid_mask] mean_emb = valid_embeddings.mean(dim=0) mean_emb = F.normalize(mean_emb, p=2, dim=0) artist_embeddings[artist_name] = mean_emb.numpy() print(f"Success: {len(artist_embeddings)}, Failed: {len(failed_artists)}") # 저장 output_path = Path(output_path) output_path.parent.mkdir(parents=True, exist_ok=True) artist_names = list(artist_embeddings.keys()) embeddings_array = np.stack([artist_embeddings[name] for name in artist_names]) np.savez_compressed( output_path, artist_names=np.array(artist_names), embeddings=embeddings_array, ) print(f"Saved: {output_path}") # 메타데이터 meta_path = output_path.with_suffix('.json') with open(meta_path, 'w') as f: json.dump({ 'num_artists': len(artist_embeddings), 'embedding_dim': config.model.embedding_dim, 'max_combinations': max_combinations, 'failed_artists': failed_artists, }, f, indent=2) print(f"Saved: {meta_path}") def load_embeddings(npz_path: str) -> Tuple[List[str], np.ndarray]: """저장된 임베딩 로드""" data = np.load(npz_path) return data['artist_names'].tolist(), data['embeddings'] def find_similar_artists( query_embedding: np.ndarray, artist_names: List[str], embeddings: np.ndarray, top_k: int = 10, ) -> List[Tuple[str, float]]: """유사 작가 검색""" query_norm = query_embedding / np.linalg.norm(query_embedding) embeddings_norm = embeddings / np.linalg.norm(embeddings, axis=1, keepdims=True) similarities = embeddings_norm @ query_norm top_indices = np.argsort(similarities)[::-1][:top_k] return [(artist_names[i], float(similarities[i])) for i in top_indices] def main(): parser = argparse.ArgumentParser(description='Three-View-Style-Embedder - Extract Embeddings') parser.add_argument('--checkpoint', type=str, required=True) parser.add_argument('--dataset_root', type=str, default='./dataset') parser.add_argument('--dataset_face_root', type=str, default='./dataset_face') parser.add_argument('--dataset_eyes_root', type=str, default='./dataset_eyes') parser.add_argument('--output', type=str, default='./embeddings/artist_embeddings.npz') parser.add_argument('--max_combinations', type=int, default=30) parser.add_argument('--batch_size', type=int, default=64) parser.add_argument('--num_workers', type=int, default=8) parser.add_argument('--device', type=str, default='cuda') parser.add_argument('--seed', type=int, default=42) args = parser.parse_args() random.seed(args.seed) np.random.seed(args.seed) torch.manual_seed(args.seed) extract_all_embeddings( checkpoint_path=args.checkpoint, dataset_root=args.dataset_root, dataset_face_root=args.dataset_face_root, dataset_eyes_root=args.dataset_eyes_root, output_path=args.output, max_combinations=args.max_combinations, batch_size=args.batch_size, num_workers=args.num_workers, device=args.device, ) if __name__ == '__main__': main()