客房场景分类器

基于DINOv2的多头酒店图像场景分类模型。

模型概述

该模型是一个从场景(Scene)、概念(Concept)、**物体(Object)**三个视角同时对酒店图像进行分类的多头深度学习模型。使用DINOv2主干网络提取强大的视觉特征,并在每个头部执行专门化的分类。

模型信息

• 模型名称: image_classifier_model_0.2
• 基础模型: facebook/dinov2-large
• 图像尺寸: 224x224
• 通道: RGB (3通道)
• 总参数: 303,252,502个 (主干网络固定)
• 可训练参数: 24,598个

分类头

Scene 头 (6个类别)

• 客房、浴室、游泳池、大堂、餐厅、其他

Concept 头 (3个类别)

• 室内、室外、特写

Object 头 (13个类别)

• 床、沙发、淋浴器、浴缸、椅子、桌子、电视、冰箱、水槽、梳妆台、镜子、其他、未分类

使用方法

使用Python调用模型

import torch
import onnxruntime as ort
import numpy as np
from PIL import Image
from torchvision import transforms
import json

# 加载模型信息
with open('image_classifier_model_0.2_model_info.json', 'r') as f:
    model_info = json.load(f)

# 加载PyTorch模型
model = torch.load('image_classifier_model_0.2.pth', map_location='cpu')
model.eval()

# 使用ONNX模型 (推理更快)
onnx_session = ort.InferenceSession('image_classifier_model_0.2.onnx')

# 图像预处理
transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])

def classify_image_pytorch(image_path):
    \"\"\"使用PyTorch模型进行图像分类\"\"\"
    image = transform(Image.open(image_path)).unsqueeze(0)
    
    with torch.no_grad():
        outputs = model(image)
        predictions = {}
        
        for head_name, logits in outputs.items():
            probabilities = torch.softmax(logits, dim=1)
            predicted_class = torch.argmax(probabilities, dim=1).item()
            confidence = probabilities[0, predicted_class].item()
            
            predictions[head_name] = {
                'class_id': predicted_class,
                'confidence': confidence,
                'probabilities': probabilities[0].tolist()
            }
    
    return predictions

def classify_image_onnx(image_path):
    \"\"\"使用ONNX模型进行图像分类 (推荐)\"\"\"
    image = transform(Image.open(image_path)).numpy()
    
    # ONNX模型推理
    input_feed = {'input': image.astype(np.float32)}
    outputs = onnx_session.run(None, input_feed)
    
    predictions = {}
    head_names = ['scene', 'concept', 'object']
    
    for i, head_name in enumerate(head_names):
        logits = outputs[i]
        probabilities = torch.softmax(torch.tensor(logits), dim=1)
        predicted_class = torch.argmax(probabilities, dim=1).item()
        confidence = probabilities[0, predicted_class].item()
        
        predictions[head_name] = {
            'class_id': predicted_class,
            'confidence': confidence,
            'probabilities': probabilities[0].tolist()
        }
    
    return predictions

# 示例使用
predictions = classify_image_onnx("hotel_room.jpg")
print("分类结果:")
for head, result in predictions.items():
    print(f"{head}: 类别 {result['class_id']}, 置信度 {result['confidence']:.4f}")

将类别ID转换为实际类别名称

def get_class_names(predictions, model_info):
    \"\"\"将类别ID转换为实际类别名称\"\"\"
    class_mappings = model_info['class_mappings']
    
    results = {}
    for head, result in predictions.items():
        class_id = result['class_id']
        if head in class_mappings:
            actual_class_id = class_mappings[head][str(class_id)]
            results[head] = {
                'class_id': actual_class_id,
                'confidence': result['confidence']
            }
    
    return results

# 类别名称转换示例
class_names = get_class_names(predictions, model_info)
print("实际类别ID:")
for head, result in class_names.items():
    print(f"{head}: {result['class_id']}")

批量处理

def classify_batch_images(image_paths):
    \"\"\"一次性分类多张图像\"\"\"
    results = []
    
    for image_path in image_paths:
        predictions = classify_image_onnx(image_path)
        results.append({
            'image_path': image_path,
            'predictions': predictions
        })
    
    return results

# 示例
image_paths = ["room1.jpg", "bathroom1.jpg", "lobby1.jpg"]
batch_results = classify_batch_images(image_paths)

for result in batch_results:
    print(f"\n图像: {result['image_path']}")
    for head, pred in result['predictions'].items():
        print(f"  {head}: 类别 {pred['class_id']}, 置信度 {pred['confidence']:.4f}")

模型文件

• image_classifier_model_0.2.pth: PyTorch模型文件
• image_classifier_model_0.2.onnx: ONNX模型文件 (推理优化)
• image_classifier_model_0.2_model_info.json: 模型元数据
• image_classifier_model_0.2_inference_example.py: 推理示例代码

模型架构

多头分类系统

输入图像 (224×224)
    ↓
DINOv2 主干网络 (冻结)
    ↓
共享特征 (1024维)
    ├─── Scene 头 → 6个类别
    ├─── Concept 头 → 3个类别
    └─── Object 头 → 13个类别

主要特点

• DINOv2主干网络: 基于强大的视觉Transformer特征提取
• 主干网络冻结: 利用预训练特征防止过拟合
• 多头设计: 3个头部进行多角度分析
• 类别权重: 自动平衡不均衡数据

预处理要求

1. 图像尺寸: 224x224像素
2. 色彩空间: RGB
3. 归一化: 使用ImageNet标准值 (mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
4. 裁剪: 中心裁剪 (center crop)
5. 支持格式: JPG、PNG、JPEG

使用案例

直接使用

• 酒店图像自动分类: 自动分类客房、浴室、大堂等场景
• 图像元数据生成: 自动提取图像的场景、概念、物体信息
• 图像数据库管理: 大量酒店图像自动标签
• 质量管理: 图像分类一致性验证

下游应用

• 酒店管理系统: 客房图像自动分类及管理
• 旅游平台: 按客房类型筛选图像
• 房地产平台: 自动提取住宿设施信息
• 图像搜索引擎: 基于多属性的图像搜索

限制

1. 领域专用: 专门针对酒店/住宿图像,在其他领域性能有限
2. 图像质量: 低质量或噪声较多的图像可能导致性能下降
3. 角度依赖性: 对特定角度拍摄的图像性能可能不同
4. 类别不平衡: 某些类别的性能可能低于其他类别

许可证

Apache 2.0 许可证

参考

该模型是Room Clusterer项目的一部分。更多信息请参考项目仓库。