src/data/preprocessing.py

import torch
import torchvision.transforms as T
from PIL import Image
import numpy as np
from typing import Dict, List, Optional, Tuple
import random


def get_transform(config: dict, mode: str = 'train') -> T.Compose:
    """获取数据预处理变换"""
    preprocessing = config.get('preprocessing', {})
    target_size = preprocessing.get('target_size', 512)
    resize_mode = preprocessing.get('resize_mode', 'center_crop')
    
    transforms_list = []
    
    # 训练和验证的不同变换
    if mode == 'train':
        # 随机裁剪
        if preprocessing.get('random_crop', True):
            transforms_list.append(T.RandomResizedCrop(
                target_size,
                scale=(0.8, 1.0),
                ratio=(0.8, 1.2)
            ))
        else:
            transforms_list.append(T.Resize(target_size, interpolation=T.InterpolationMode.BILINEAR))
        
        # 随机水平翻转
        if preprocessing.get('random_flip', True):
            transforms_list.append(T.RandomHorizontalFlip(p=0.5))
        
        # 颜色抖动
        if preprocessing.get('color_jitter', 0.05) > 0:
            color_jitter = preprocessing['color_jitter']
            transforms_list.append(T.ColorJitter(
                brightness=color_jitter,
                contrast=color_jitter,
                saturation=color_jitter,
                hue=min(0.1, color_jitter)
            ))
        
        # 随机旋转
        if preprocessing.get('random_rotation', 0.0) > 0:
            max_angle = preprocessing['random_rotation']
            transforms_list.append(T.RandomRotation(degrees=(-max_angle, max_angle)))
    
    else:  # 验证/测试模式
        # 中心裁剪
        if resize_mode == 'center_crop':
            transforms_list.extend([
                T.Resize(target_size, interpolation=T.InterpolationMode.BILINEAR),
                T.CenterCrop(target_size)
            ])
        elif resize_mode == 'resize':
            transforms_list.append(T.Resize(target_size, interpolation=T.InterpolationMode.BILINEAR))
        elif resize_mode == 'random_crop':
            transforms_list.append(T.RandomCrop(target_size))
        else:
            raise ValueError(f"未知的resize_mode: {resize_mode}")
    
    # 转换为Tensor
    transforms_list.append(T.ToTensor())
    
    # 归一化
    normalize_config = preprocessing.get('normalize', {})
    mean = normalize_config.get('mean', [0.5, 0.5, 0.5])
    std = normalize_config.get('std', [0.5, 0.5, 0.5])
    transforms_list.append(T.Normalize(mean=mean, std=std))
    
    return T.Compose(transforms_list)


class TextPreprocessor:
    """文本预处理器"""
    def __init__(self, config: dict):
        self.config = config.get('preprocessing', {})
        
        # 文本处理参数
        self.max_length = self.config.get('max_length', 77)
        self.truncation = self.config.get('truncation', True)
        self.padding = self.config.get('padding', 'max_length')
        
        # 尝试加载tokenizer
        self.tokenizer = None
        self._init_tokenizer()
    
    def _init_tokenizer(self):
        """初始化tokenizer"""
        try:
            from transformers import CLIPTokenizer
            tokenizer_name = self.config.get('tokenizer', 'openai/clip-vit-base-patch32')
            self.tokenizer = CLIPTokenizer.from_pretrained(tokenizer_name)
        except ImportError:
            print("警告: 未安装transformers，无法使用CLIP tokenizer")
        except Exception as e:
            print(f"加载tokenizer失败: {e}")
    
    def preprocess_text(self, text: str) -> Dict:
        """预处理文本"""
        if self.tokenizer is not None:
            # 使用CLIP tokenizer
            inputs = self.tokenizer(
                text,
                max_length=self.max_length,
                padding=self.padding,
                truncation=self.truncation,
                return_tensors="pt"
            )
            return {
                'input_ids': inputs['input_ids'].squeeze(0),
                'attention_mask': inputs['attention_mask'].squeeze(0)
            }
        else:
            # 简单的文本处理
            return {
                'text': text,
                'length': len(text)
            }
    
    def batch_preprocess(self, texts: List[str]) -> Dict:
        """批量预处理文本"""
        if self.tokenizer is not None:
            inputs = self.tokenizer(
                texts,
                max_length=self.max_length,
                padding=self.padding,
                truncation=self.truncation,
                return_tensors="pt"
            )
            return inputs
        else:
            return {'texts': texts}


class ImagePreprocessor:
    """图像预处理器"""
    def __init__(self, config: dict):
        self.config = config.get('preprocessing', {})
        self.transform = get_transform(config, mode='train')
    
    def preprocess_image(self, image: Image.Image) -> torch.Tensor:
        """预处理单张图像"""
        return self.transform(image)
    
    def batch_preprocess(self, images: List[Image.Image]) -> torch.Tensor:
        """批量预处理图像"""
        return torch.stack([self.transform(img) for img in images])
    
    def preprocess_for_vae(self, image: torch.Tensor) -> torch.Tensor:
        """为VAE编码预处理图像"""
        # VAE期望输入在[-1, 1]范围内
        return image * 2.0 - 1.0
    
    def postprocess_from_vae(self, latents: torch.Tensor) -> torch.Tensor:
        """从VAE解码后处理图像"""
        # 将[-1, 1]范围转换回[0, 1]
        return (latents + 1.0) / 2.0


class DataPreprocessor:
    """数据预处理器（整合文本和图像处理）"""
    def __init__(self, config: dict):
        self.config = config
        self.image_preprocessor = ImagePreprocessor(config)
        self.text_preprocessor = TextPreprocessor(config)
        
        # 文本编码器
        self.text_encoder = None
        self._init_text_encoder()
    
    def _init_text_encoder(self):
        """初始化文本编码器"""
        try:
            from transformers import CLIPTextModel
            model_name = self.config.get('preprocessing', {}).get('text_encoder', 'openai/clip-vit-base-patch32')
            self.text_encoder = CLIPTextModel.from_pretrained(model_name)
            
            # 冻结参数
            for param in self.text_encoder.parameters():
                param.requires_grad = False
            
            # 设置为评估模式
            self.text_encoder.eval()
            
            print(f"已加载文本编码器: {model_name}")
        except Exception as e:
            print(f"加载文本编码器失败: {e}")
    
    def encode_text(self, text: str) -> torch.Tensor:
        """编码文本为嵌入向量"""
        if self.text_encoder is None:
            raise ValueError("文本编码器未初始化")
        
        # 预处理文本
        inputs = self.text_preprocessor.preprocess_text(text)
        
        # 编码
        with torch.no_grad():
            if 'input_ids' in inputs:
                outputs = self.text_encoder(
                    input_ids=inputs['input_ids'].unsqueeze(0),
                    attention_mask=inputs['attention_mask'].unsqueeze(0) if 'attention_mask' in inputs else None
                )
                return outputs.last_hidden_state.squeeze(0)
            else:
                # 回退到简单的嵌入
                return torch.randn(77, 768)  # 默认维度
    
    def batch_encode_text(self, texts: List[str]) -> torch.Tensor:
        """批量编码文本"""
        if self.text_encoder is None:
            raise ValueError("文本编码器未初始化")
        
        # 预处理文本
        inputs = self.text_preprocessor.batch_preprocess(texts)
        
        # 编码
        with torch.no_grad():
            if 'input_ids' in inputs:
                outputs = self.text_encoder(
                    input_ids=inputs['input_ids'],
                    attention_mask=inputs.get('attention_mask', None)
                )
                return outputs.last_hidden_state
            else:
                # 回退到简单的嵌入
                batch_size = len(texts)
                return torch.randn(batch_size, 77, 768)
    
    def preprocess_batch(self, batch: List[Dict]) -> Dict:
        """预处理批次数据"""
        images = [item['image'] for item in batch]
        texts = [item['text'] for item in batch]
        
        # 预处理图像
        image_tensors = self.image_preprocessor.batch_preprocess(images)
        
        # 编码文本
        if self.text_encoder is not None:
            text_embeddings = self.batch_encode_text(texts)
        else:
            text_embeddings = None
        
        return {
            'images': image_tensors,
            'text_embeddings': text_embeddings,
            'texts': texts,
            'image_paths': [item.get('image_path', '') for item in batch]
        }


def test_preprocessing():
    """测试预处理"""
    import yaml
    from PIL import Image
    
    # 创建测试图像
    test_image = Image.new('RGB', (512, 512), color='red')
    
    # 加载配置
    with open('configs/data/laion_filtered.yaml', 'r') as f:
        config = yaml.safe_load(f)
    
    # 测试图像预处理
    image_preprocessor = ImagePreprocessor(config)
    processed_image = image_preprocessor.preprocess_image(test_image)
    
    print(f"原始图像: {test_image.size}")
    print(f"处理后图像形状: {processed_image.shape}")
    
    # 测试文本预处理
    text_preprocessor = TextPreprocessor(config)
    processed_text = text_preprocessor.preprocess_text("A red square image")
    
    print(f"文本处理结果: {processed_text}")
    
    # 测试数据预处理器
    data_preprocessor = DataPreprocessor(config)
    
    test_batch = [
        {'image': test_image, 'text': "A red square"},
        {'image': test_image, 'text': "A blue circle"}
    ]
    
    processed_batch = data_preprocessor.preprocess_batch(test_batch)
    
    print(f"批次图像形状: {processed_batch['images'].shape}")
    print(f"文本嵌入形状: {processed_batch['text_embeddings'].shape if processed_batch['text_embeddings'] is not None else 'None'}")
    
    return processed_batch


if __name__ == '__main__':
    test_preprocessing()