ondame/image-classifier

1---2library_name: transformers3tags:4- image-classification5- multi-head-classification6- room-classification7- dinov28- computer-vision9- scene-classification10license: apache-2.011language:12- en13pipeline_tag: image-classification14base_model:15- facebook/dinov2-large16---17 18# Room Scene Classifier19 20DINOv2 기반 멀티헤드 호텔 이미지 장면 분류 모델입니다.21 22## 모델 개요23 24이 모델은 호텔 이미지를 **Scene(장면)**, **Concept(개념)**, **Object(객체)** 3가지 관점으로 동시에 분류하는 멀티헤드 딥러닝 모델입니다. DINOv2 백본을 사용하여 강력한 비전 특징을 추출하고, 각 헤드에서 특화된 분류를 수행합니다.25 26## 모델 정보27 28- **모델명**: `image_classifier_model_0.2`29- **기반 모델**: `facebook/dinov2-large`30- **이미지 크기**: 224x22431- **채널**: RGB (3채널)32- **총 파라미터**: 303,252,502개 (백본 고정)33- **훈련 가능 파라미터**: 24,598개34 35## 분류 헤드36 37### Scene 헤드 (6개 클래스)38- 객실, 욕실, 수영장, 로비, 레스토랑, 기타39 40### Concept 헤드 (3개 클래스)  41- 실내, 야외, 클로즈업42 43### Object 헤드 (13개 클래스)44- 침대, 소파, 샤워기, 욕조, 의자, 테이블, TV, 냉장고, 싱크대, 화장대, 거울, 기타, 미분류45 46## 사용법47 48### Python으로 모델 사용49 50```python51import torch52import onnxruntime as ort53import numpy as np54from PIL import Image55from torchvision import transforms56import json57 58# 모델 정보 로드59with open('image_classifier_model_0.2_model_info.json', 'r') as f:60    model_info = json.load(f)61 62# PyTorch 모델 로드63model = torch.load('image_classifier_model_0.2.pth', map_location='cpu')64model.eval()65 66# ONNX 모델 사용 (더 빠른 추론)67onnx_session = ort.InferenceSession('image_classifier_model_0.2.onnx')68 69# 이미지 전처리70transform = transforms.Compose([71    transforms.Resize((224, 224)),72    transforms.CenterCrop(224),73    transforms.ToTensor(),74    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])75])76 77def classify_image_pytorch(image_path):78    """PyTorch 모델을 사용한 이미지 분류"""79    image = transform(Image.open(image_path)).unsqueeze(0)80    81    with torch.no_grad():82        outputs = model(image)83        predictions = {}84        85        for head_name, logits in outputs.items():86            probabilities = torch.softmax(logits, dim=1)87            predicted_class = torch.argmax(probabilities, dim=1).item()88            confidence = probabilities[0, predicted_class].item()89            90            predictions[head_name] = {91                'class_id': predicted_class,92                'confidence': confidence,93                'probabilities': probabilities[0].tolist()94            }95    96    return predictions97 98def classify_image_onnx(image_path):99    """ONNX 모델을 사용한 이미지 분류 (권장)"""100    image = transform(Image.open(image_path)).numpy()101    102    # ONNX 모델 추론103    input_feed = {'input': image.astype(np.float32)}104    outputs = onnx_session.run(None, input_feed)105    106    predictions = {}107    head_names = ['scene', 'concept', 'object']108    109    for i, head_name in enumerate(head_names):110        logits = outputs[i]111        probabilities = torch.softmax(torch.tensor(logits), dim=1)112        predicted_class = torch.argmax(probabilities, dim=1).item()113        confidence = probabilities[0, predicted_class].item()114        115        predictions[head_name] = {116            'class_id': predicted_class,117            'confidence': confidence,118            'probabilities': probabilities[0].tolist()119        }120    121    return predictions122 123# 예시 사용124predictions = classify_image_onnx("hotel_room.jpg")125print("분류 결과:")126for head, result in predictions.items():127    print(f"{head}: 클래스 {result['class_id']}, 신뢰도 {result['confidence']:.4f}")128```129 130### 클래스 ID를 실제 클래스명으로 변환131 132```python133def get_class_names(predictions, model_info):134    """클래스 ID를 실제 클래스명으로 변환"""135    class_mappings = model_info['class_mappings']136    137    results = {}138    for head, result in predictions.items():139        class_id = result['class_id']140        if head in class_mappings:141            actual_class_id = class_mappings[head][str(class_id)]142            results[head] = {143                'class_id': actual_class_id,144                'confidence': result['confidence']145            }146    147    return results148 149# 클래스명 변환 예시150class_names = get_class_names(predictions, model_info)151print("실제 클래스 ID:")152for head, result in class_names.items():153    print(f"{head}: {result['class_id']}")154```155 156### 배치 처리157 158```python159def classify_batch_images(image_paths):160    """여러 이미지를 한 번에 분류"""161    results = []162    163    for image_path in image_paths:164        predictions = classify_image_onnx(image_path)165        results.append({166            'image_path': image_path,167            'predictions': predictions168        })169    170    return results171 172# 예시173image_paths = ["room1.jpg", "bathroom1.jpg", "lobby1.jpg"]174batch_results = classify_batch_images(image_paths)175 176for result in batch_results:177    print(f"\n이미지: {result['image_path']}")178    for head, pred in result['predictions'].items():179        print(f"  {head}: 클래스 {pred['class_id']}, 신뢰도 {pred['confidence']:.4f}")180```181 182## 모델 파일183 184- `image_classifier_model_0.2.pth`: PyTorch 모델 파일185- `image_classifier_model_0.2.onnx`: ONNX 모델 파일 (추론 최적화)186- `image_classifier_model_0.2_model_info.json`: 모델 메타데이터187- `image_classifier_model_0.2_inference_example.py`: 추론 예제 코드188 189 190## 모델 아키텍처191 192### 멀티헤드 분류 시스템193```194입력 이미지 (224×224)195    ↓196DINOv2 백본 (Frozen)197    ↓198공통 특징 (1024차원)199    ├─── Scene 헤드 → 6개 클래스200    ├─── Concept 헤드 → 3개 클래스201    └─── Object 헤드 → 13개 클래스202```203 204### 주요 특징205- **DINOv2 백본**: 강력한 비전 트랜스포머 기반 특징 추출206- **백본 고정**: 사전훈련된 특징을 활용하여 과적합 방지207- **멀티헤드**: 3개 헤드로 다각도 분석208- **클래스 가중치**: 불균형 데이터 자동 보정209 210## 전처리 요구사항211 2121. **이미지 크기**: 224x224 픽셀2132. **색상 공간**: RGB2143. **정규화**: ImageNet 표준값 사용 (mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])2154. **크롭**: 중앙 크롭 (center crop)2165. **지원 형식**: JPG, PNG, JPEG217 218## 사용 사례219 220### 직접 사용221 222- **호텔 이미지 자동 분류**: 객실, 욕실, 로비 등 장면별 자동 분류223- **이미지 메타데이터 생성**: 이미지의 장면, 개념, 객체 정보 자동 추출224- **이미지 데이터베이스 관리**: 대량의 호텔 이미지 자동 태깅225- **품질 관리**: 이미지 분류 일관성 검증226 227### 다운스트림 사용228 229- **호텔 관리 시스템**: 객실 이미지 자동 분류 및 관리230- **여행 플랫폼**: 객실 타입별 이미지 필터링231- **부동산 플랫폼**: 숙소 시설 정보 자동 추출232- **이미지 검색 엔진**: 다중 속성 기반 이미지 검색233 234## 제한사항235 2361. **도메인 특화**: 호텔/숙소 이미지에 특화되어 있어 다른 도메인에서는 성능이 제한적입니다.2372. **이미지 품질**: 저화질이나 노이즈가 많은 이미지에서는 성능이 저하될 수 있습니다.2383. **각도 의존성**: 특정 각도에서 촬영된 이미지에 대해 성능이 다를 수 있습니다.2394. **클래스 불균형**: 일부 클래스는 다른 클래스보다 성능이 낮을 수 있습니다.240 241## 라이선스242 243Apache 2.0 License244 245## 참고246 247이 모델은 Room Clusterer 프로젝트의 일부로 개발되었습니다. 더 자세한 정보는 [프로젝트 저장소](https://github.com/tportio/content-ml-trainer)를 참조하세요.