| 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}")
|
|
|
|
|
| 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')
|
|
|
|
|
| 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:
|
|
|
| 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编码预处理图像"""
|
|
|
| return image * 2.0 - 1.0
|
|
|
| def postprocess_from_vae(self, latents: torch.Tensor) -> torch.Tensor:
|
| """从VAE解码后处理图像"""
|
|
|
| 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() |
