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()