src/data/dataset.py

import torch
from torch.utils.data import Dataset, DataLoader
from PIL import Image
import pandas as pd
import os
from typing import Dict, List, Optional, Tuple
import numpy as np
import json
from pathlib import Path


class LAIONDataset(Dataset):
    """LAION数据集"""
    def __init__(self, config: dict, transform=None, split: str = 'train'):
        self.config = config
        self.transform = transform
        self.split = split
        
        # 加载元数据
        metadata_path = config['dataset'].get('metadata_file', './data/laion/metadata.parquet')
        self.metadata = pd.read_parquet(metadata_path)
        
        # 应用过滤条件
        self._apply_filters(config.get('filters', {}))
        
        # 数据拆分
        self._split_data(config.get('split', {}))
        
        # 缓存
        self.use_cache = config.get('use_cache', False)
        self.cache_dir = config.get('cache_dir', './data/cache')
        if self.use_cache:
            os.makedirs(self.cache_dir, exist_ok=True)
        
        # 限制样本数
        max_samples = config.get('max_samples', None)
        if max_samples is not None and len(self.metadata) > max_samples:
            self.metadata = self.metadata.sample(max_samples, random_state=42)
        
        # 文本缓存
        self.text_cache = {}
        
        print(f"数据集加载完成: {len(self.metadata)} 个样本 ({split}集)")
    
    def _apply_filters(self, filters: dict):
        """应用过滤条件"""
        if 'aesthetic_score' in filters:
            threshold = filters['aesthetic_score']
            if 'aesthetic_score' in self.metadata.columns:
                self.metadata = self.metadata[self.metadata['aesthetic_score'] >= threshold]
        
        if 'watermark_prob' in filters:
            threshold = filters['watermark_prob']
            if 'watermark_prob' in self.metadata.columns:
                self.metadata = self.metadata[self.metadata['watermark_prob'] <= threshold]
        
        if 'nsfw' in filters and not filters['nsfw']:
            if 'NSFW' in self.metadata.columns:
                self.metadata = self.metadata[self.metadata['NSFW'] != 'NSFW']
    
    def _split_data(self, split_config: dict):
        """拆分数据集"""
        if self.split not in ['train', 'val']:
            return
        
        train_ratio = split_config.get('train', 0.95)
        val_ratio = split_config.get('val', 0.05)
        
        # 确保拆分比例之和为1
        total = train_ratio + val_ratio
        train_ratio /= total
        val_ratio /= total
        
        # 随机拆分
        seed = split_config.get('seed', 42)
        shuffled = self.metadata.sample(frac=1, random_state=seed).reset_index(drop=True)
        
        if self.split == 'train':
            split_point = int(len(shuffled) * train_ratio)
            self.metadata = shuffled[:split_point]
        else:
            split_point = int(len(shuffled) * train_ratio)
            self.metadata = shuffled[split_point:]
    
    def __len__(self) -> int:
        return len(self.metadata)
    
    def _get_image_path(self, row) -> str:
        """获取图像路径"""
        # 尝试不同的列名
        for col in ['image_file', 'filepath', 'path', 'url_local']:
            if col in row:
                path = row[col]
                # 如果是相对路径，添加基础路径
                if not os.path.isabs(path):
                    base_path = self.config['dataset'].get('path', './data/laion')
                    path = os.path.join(base_path, path)
                return path
        
        # 如果没有找到路径，使用URL哈希
        if 'url' in row:
            import hashlib
            url_hash = hashlib.md5(row['url'].encode()).hexdigest()
            base_path = self.config['dataset'].get('path', './data/laion')
            path = os.path.join(base_path, f"{url_hash}.jpg")
            return path
        
        raise ValueError(f"无法找到图像路径: {row}")
    
    def __getitem__(self, idx: int) -> Dict:
        row = self.metadata.iloc[idx]
        
        # 缓存键
        cache_key = f"{self.split}_{idx}"
        
        # 检查缓存
        if self.use_cache and cache_key in self.text_cache:
            text_embedding = self.text_cache[cache_key]
        else:
            # 获取文本描述
            text = row.get('caption', row.get('text', row.get('description', '')))
            
            # 这里应该调用文本编码器，但为了简化，我们返回原始文本
            # 在实际使用中，应该使用预训练的CLIP编码器
            text_embedding = text
            
            # 缓存
            if self.use_cache:
                self.text_cache[cache_key] = text_embedding
        
        # 获取图像
        try:
            image_path = self._get_image_path(row)
            image = Image.open(image_path).convert('RGB')
            
            # 应用变换
            if self.transform:
                image = self.transform(image)
        except Exception as e:
            # 如果图像加载失败，返回一个空白图像
            print(f"加载图像失败 {image_path}: {e}")
            image = torch.zeros(3, 512, 512)
            text = "invalid image"
            text_embedding = text
        
        return {
            'image': image,
            'text': text_embedding if isinstance(text_embedding, str) else '',
            'text_embedding': text_embedding if not isinstance(text_embedding, str) else None,
            'image_path': image_path if 'image_path' in locals() else '',
            'index': idx
        }


class TextImageDataset(Dataset):
    """文本-图像对数据集"""
    def __init__(self, image_dir: str, caption_file: str, transform=None):
        self.image_dir = image_dir
        self.transform = transform
        
        # 加载标注文件
        if caption_file.endswith('.json'):
            with open(caption_file, 'r') as f:
                self.captions = json.load(f)
        elif caption_file.endswith('.csv'):
            self.captions = pd.read_csv(caption_file)
        else:
            raise ValueError(f"不支持的标注文件格式: {caption_file}")
        
        # 验证图像文件是否存在
        self.valid_samples = []
        for item in self.captions:
            if isinstance(item, dict):
                image_name = item.get('image_name', item.get('file_name', ''))
                caption = item.get('caption', '')
            else:
                image_name = item[0]
                caption = item[1]
            
            image_path = os.path.join(self.image_dir, image_name)
            if os.path.exists(image_path):
                self.valid_samples.append((image_path, caption))
        
        print(f"找到 {len(self.valid_samples)} 个有效样本")
    
    def __len__(self) -> int:
        return len(self.valid_samples)
    
    def __getitem__(self, idx: int) -> Dict:
        image_path, caption = self.valid_samples[idx]
        
        # 加载图像
        image = Image.open(image_path).convert('RGB')
        
        # 应用变换
        if self.transform:
            image = self.transform(image)
        
        return {
            'image': image,
            'text': caption,
            'image_path': image_path
        }


class CachedDataset(Dataset):
    """缓存数据集，加速训练"""
    def __init__(self, dataset: Dataset, cache_dir: str = './cache'):
        self.dataset = dataset
        self.cache_dir = cache_dir
        os.makedirs(cache_dir, exist_ok=True)
        
        self.cache_files = []
        for i in range(len(dataset)):
            cache_file = os.path.join(cache_dir, f'sample_{i}.pt')
            self.cache_files.append(cache_file)
    
    def __len__(self) -> int:
        return len(self.dataset)
    
    def __getitem__(self, idx: int) -> Dict:
        cache_file = self.cache_files[idx]
        
        # 如果缓存存在，直接加载
        if os.path.exists(cache_file):
            try:
                return torch.load(cache_file)
            except:
                pass
        
        # 否则，从原始数据集加载并缓存
        sample = self.dataset[idx]
        torch.save(sample, cache_file)
        
        return sample


def create_data_loaders(config: dict) -> Tuple[DataLoader, Optional[DataLoader]]:
    """创建数据加载器"""
    from .preprocessing import get_transform
    
    # 获取数据变换
    train_transform = get_transform(config, mode='train')
    val_transform = get_transform(config, mode='val')
    
    # 创建数据集
    train_dataset = LAIONDataset(config, transform=train_transform, split='train')
    val_dataset = LAIONDataset(config, transform=val_transform, split='val')
    
    # 可选：启用缓存
    if config.get('cache_dataset', True):
        train_dataset = CachedDataset(train_dataset, cache_dir='./data/cache/train')
        val_dataset = CachedDataset(val_dataset, cache_dir='./data/cache/val')
    
    # 创建数据加载器
    train_loader = DataLoader(
        train_dataset,
        batch_size=config.get('batch_size', 1),
        shuffle=config.get('shuffle', True),
        num_workers=config.get('num_workers', 2),
        pin_memory=config.get('pin_memory', True),
        prefetch_factor=config.get('prefetch_factor', 2),
        persistent_workers=config.get('persistent_workers', True)
    )
    
    val_loader = DataLoader(
        val_dataset,
        batch_size=1,  # 验证时批次大小为1
        shuffle=False,
        num_workers=config.get('num_workers', 2),
        pin_memory=True
    )
    
    return train_loader, val_loader


def test_dataset():
    """测试数据集"""
    import yaml
    
    # 加载配置
    with open('configs/data/laion_filtered.yaml', 'r') as f:
        config = yaml.safe_load(f)
    
    # 创建数据集
    dataset = LAIONDataset(config, split='train')
    
    # 测试样本
    sample = dataset[0]
    print(f"样本键: {list(sample.keys())}")
    print(f"图像形状: {sample['image'].shape if hasattr(sample['image'], 'shape') else type(sample['image'])}")
    print(f"文本: {sample['text'][:100]}...")
    
    return dataset


if __name__ == '__main__':
    test_dataset()