wehub-resource-sync
316 lines
11 KiB
Python
316 lines
11 KiB
Python
import json
|
||
import os
|
||
from pathlib import Path
|
||
|
||
import datasets
|
||
import lmdb
|
||
import numpy as np
|
||
import torch
|
||
from PIL import Image
|
||
from torch.utils.data import Dataset
|
||
|
||
from .lmdb import get_array_shape_from_lmdb, retrieve_row_from_lmdb
|
||
|
||
|
||
class TextDataset(Dataset):
|
||
def __init__(self, prompt_path, extended_prompt_path=None):
|
||
with open(prompt_path, encoding="utf-8") as f:
|
||
self.prompt_list = [line.rstrip() for line in f]
|
||
|
||
if extended_prompt_path is not None:
|
||
with open(extended_prompt_path, encoding="utf-8") as f:
|
||
self.extended_prompt_list = [line.rstrip() for line in f]
|
||
assert len(self.extended_prompt_list) == len(self.prompt_list)
|
||
else:
|
||
self.extended_prompt_list = None
|
||
|
||
def __len__(self):
|
||
return len(self.prompt_list)
|
||
|
||
def __getitem__(self, idx):
|
||
batch = {
|
||
"prompts": self.prompt_list[idx],
|
||
"idx": idx,
|
||
}
|
||
if self.extended_prompt_list is not None:
|
||
batch["extended_prompts"] = self.extended_prompt_list[idx]
|
||
return batch
|
||
|
||
|
||
class TwoTextDataset(Dataset):
|
||
"""Dataset that returns two text prompts per sample for prompt-switch training.
|
||
|
||
The dataset behaves similarly to :class:`TextDataset` but instead of a single
|
||
prompt, it provides *two* prompts – typically the first prompt is used for the
|
||
first segment of the video, and the second prompt is used after a temporal
|
||
switch during training.
|
||
|
||
Args:
|
||
prompt_path (str): Path to a text file containing the *first* prompt for
|
||
each sample. One prompt per line.
|
||
switch_prompt_path (str): Path to a text file containing the *second*
|
||
prompt for each sample. Must have the **same number of lines** as
|
||
``prompt_path`` so that prompts are paired 1-to-1.
|
||
"""
|
||
|
||
def __init__(self, prompt_path: str, switch_prompt_path: str):
|
||
# Load the first-segment prompts.
|
||
with open(prompt_path, encoding="utf-8") as f:
|
||
self.prompt_list = [line.rstrip() for line in f]
|
||
|
||
# Load the second-segment prompts.
|
||
with open(switch_prompt_path, encoding="utf-8") as f:
|
||
self.switch_prompt_list = [line.rstrip() for line in f]
|
||
|
||
assert len(self.switch_prompt_list) == len(self.prompt_list), (
|
||
"The two prompt files must contain the same number of lines so that "
|
||
"each first-segment prompt is paired with exactly one second-segment prompt."
|
||
)
|
||
|
||
def __len__(self):
|
||
return len(self.prompt_list)
|
||
|
||
def __getitem__(self, idx):
|
||
return {
|
||
"prompts": self.prompt_list[idx], # first-segment prompt
|
||
"switch_prompts": self.switch_prompt_list[idx], # second-segment prompt
|
||
"idx": idx,
|
||
}
|
||
|
||
|
||
class ODERegressionLMDBDataset(Dataset):
|
||
def __init__(self, data_path: str, max_pair: int = int(1e8)):
|
||
self.env = lmdb.open(data_path, readonly=True, lock=False, readahead=False, meminit=False)
|
||
|
||
self.latents_shape = get_array_shape_from_lmdb(self.env, "latents")
|
||
self.max_pair = max_pair
|
||
|
||
def __len__(self):
|
||
return min(self.latents_shape[0], self.max_pair)
|
||
|
||
def __getitem__(self, idx):
|
||
"""
|
||
Outputs:
|
||
- prompts: List of Strings
|
||
- latents: Tensor of shape (num_denoising_steps, num_frames, num_channels, height, width). It is ordered from pure noise to clean image.
|
||
"""
|
||
latents = retrieve_row_from_lmdb(self.env, "latents", np.float16, idx, shape=self.latents_shape[1:])
|
||
|
||
if len(latents.shape) == 4:
|
||
latents = latents[None, ...]
|
||
|
||
prompts = retrieve_row_from_lmdb(self.env, "prompts", str, idx)
|
||
return {"prompts": prompts, "ode_latent": torch.tensor(latents, dtype=torch.float32)}
|
||
|
||
|
||
class ShardingLMDBDataset(Dataset):
|
||
def __init__(self, data_path: str, max_pair: int = int(1e8)):
|
||
self.envs = []
|
||
self.index = []
|
||
|
||
for fname in sorted(os.listdir(data_path)):
|
||
path = os.path.join(data_path, fname)
|
||
env = lmdb.open(path, readonly=True, lock=False, readahead=False, meminit=False)
|
||
self.envs.append(env)
|
||
|
||
self.latents_shape = [None] * len(self.envs)
|
||
for shard_id, env in enumerate(self.envs):
|
||
self.latents_shape[shard_id] = get_array_shape_from_lmdb(env, "latents")
|
||
for local_i in range(self.latents_shape[shard_id][0]):
|
||
self.index.append((shard_id, local_i))
|
||
|
||
self.max_pair = max_pair
|
||
|
||
def __len__(self):
|
||
return len(self.index)
|
||
|
||
def __getitem__(self, idx):
|
||
"""
|
||
Outputs:
|
||
- prompts: List of Strings
|
||
- latents: Tensor of shape (num_denoising_steps, num_frames, num_channels, height, width). It is ordered from pure noise to clean image.
|
||
"""
|
||
shard_id, local_idx = self.index[idx]
|
||
|
||
latents = retrieve_row_from_lmdb(
|
||
self.envs[shard_id], "latents", np.float16, local_idx, shape=self.latents_shape[shard_id][1:]
|
||
)
|
||
|
||
if len(latents.shape) == 4:
|
||
latents = latents[None, ...]
|
||
|
||
prompts = retrieve_row_from_lmdb(self.envs[shard_id], "prompts", str, local_idx)
|
||
|
||
return {"prompts": prompts, "ode_latent": torch.tensor(latents, dtype=torch.float32)}
|
||
|
||
|
||
class TextImagePairDataset(Dataset):
|
||
def __init__(self, data_dir, transform=None, eval_first_n=-1, pad_to_multiple_of=None):
|
||
"""
|
||
Args:
|
||
data_dir (str): Path to the directory containing:
|
||
- target_crop_info_*.json (metadata file)
|
||
- */ (subdirectory containing images with matching aspect ratio)
|
||
transform (callable, optional): Optional transform to be applied on the image
|
||
"""
|
||
self.transform = transform
|
||
data_dir = Path(data_dir)
|
||
|
||
# Find the metadata JSON file
|
||
metadata_files = list(data_dir.glob("target_crop_info_*.json"))
|
||
if not metadata_files:
|
||
raise FileNotFoundError(f"No metadata file found in {data_dir}")
|
||
if len(metadata_files) > 1:
|
||
raise ValueError(f"Multiple metadata files found in {data_dir}")
|
||
|
||
metadata_path = metadata_files[0]
|
||
# Extract aspect ratio from metadata filename (e.g. target_crop_info_26-15.json -> 26-15)
|
||
aspect_ratio = metadata_path.stem.split("_")[-1]
|
||
|
||
# Use aspect ratio subfolder for images
|
||
self.image_dir = data_dir / aspect_ratio
|
||
if not self.image_dir.exists():
|
||
raise FileNotFoundError(f"Image directory not found: {self.image_dir}")
|
||
|
||
# Load metadata
|
||
with open(metadata_path) as f:
|
||
self.metadata = json.load(f)
|
||
|
||
eval_first_n = eval_first_n if eval_first_n != -1 else len(self.metadata)
|
||
self.metadata = self.metadata[:eval_first_n]
|
||
|
||
# Verify all images exist
|
||
for item in self.metadata:
|
||
image_path = self.image_dir / item["file_name"]
|
||
if not image_path.exists():
|
||
raise FileNotFoundError(f"Image not found: {image_path}")
|
||
|
||
self.dummy_prompt = "DUMMY PROMPT"
|
||
self.pre_pad_len = len(self.metadata)
|
||
if pad_to_multiple_of is not None and len(self.metadata) % pad_to_multiple_of != 0:
|
||
# Duplicate the last entry
|
||
self.metadata += [self.metadata[-1]] * (pad_to_multiple_of - len(self.metadata) % pad_to_multiple_of)
|
||
|
||
def __len__(self):
|
||
return len(self.metadata)
|
||
|
||
def __getitem__(self, idx):
|
||
"""
|
||
Returns:
|
||
dict: A dictionary containing:
|
||
- image: PIL Image
|
||
- caption: str
|
||
- target_bbox: list of int [x1, y1, x2, y2]
|
||
- target_ratio: str
|
||
- type: str
|
||
- origin_size: tuple of int (width, height)
|
||
"""
|
||
item = self.metadata[idx]
|
||
|
||
# Load image
|
||
image_path = self.image_dir / item["file_name"]
|
||
image = Image.open(image_path).convert("RGB")
|
||
|
||
# Apply transform if specified
|
||
if self.transform:
|
||
image = self.transform(image)
|
||
|
||
return {
|
||
"image": image,
|
||
"prompts": item["caption"],
|
||
"target_bbox": item["target_crop"]["target_bbox"],
|
||
"target_ratio": item["target_crop"]["target_ratio"],
|
||
"type": item["type"],
|
||
"origin_size": (item["origin_width"], item["origin_height"]),
|
||
"idx": idx,
|
||
}
|
||
|
||
|
||
class TwoTextDataset(Dataset):
|
||
"""Dataset that returns two text prompts per sample for prompt-switch training.
|
||
|
||
The dataset behaves similarly to :class:`TextDataset` but instead of a single
|
||
prompt, it provides *two* prompts – typically the first prompt is used for the
|
||
first segment of the video, and the second prompt is used after a temporal
|
||
switch during training.
|
||
|
||
Args:
|
||
prompt_path (str): Path to a text file containing the *first* prompt for
|
||
each sample. One prompt per line.
|
||
switch_prompt_path (str): Path to a text file containing the *second*
|
||
prompt for each sample. Must have the **same number of lines** as
|
||
``prompt_path`` so that prompts are paired 1-to-1.
|
||
"""
|
||
|
||
def __init__(self, prompt_path: str, switch_prompt_path: str):
|
||
# Load the first-segment prompts.
|
||
with open(prompt_path, encoding="utf-8") as f:
|
||
self.prompt_list = [line.rstrip() for line in f]
|
||
|
||
# Load the second-segment prompts.
|
||
with open(switch_prompt_path, encoding="utf-8") as f:
|
||
self.switch_prompt_list = [line.rstrip() for line in f]
|
||
|
||
assert len(self.switch_prompt_list) == len(self.prompt_list), (
|
||
"The two prompt files must contain the same number of lines so that "
|
||
"each first-segment prompt is paired with exactly one second-segment prompt."
|
||
)
|
||
|
||
def __len__(self):
|
||
return len(self.prompt_list)
|
||
|
||
def __getitem__(self, idx):
|
||
return {
|
||
"prompts": self.prompt_list[idx], # first-segment prompt
|
||
"switch_prompts": self.switch_prompt_list[idx], # second-segment prompt
|
||
"idx": idx,
|
||
}
|
||
|
||
|
||
class MultiTextDataset(Dataset):
|
||
"""Dataset for multiple‑segment prompts stored in a JSONL file.
|
||
|
||
Each line is a JSON object, e.g.
|
||
{"prompts": ["a cat", "a dog", "a bird"]}
|
||
|
||
Args:
|
||
prompt_path (str): JSONL file path
|
||
field (str): Field name to save the list of strings, default "prompts"
|
||
cache_dir (str | None): Cache directory for HF Datasets
|
||
"""
|
||
|
||
def __init__(self, prompt_path: str, field: str = "prompts", cache_dir: str | None = None):
|
||
self.ds = datasets.load_dataset(
|
||
"json",
|
||
data_files=prompt_path,
|
||
split="train",
|
||
cache_dir=cache_dir,
|
||
streaming=False,
|
||
)
|
||
|
||
assert len(self.ds) > 0, "JSONL is empty"
|
||
assert field in self.ds.column_names, f"Field '{field}' is missing"
|
||
|
||
# Check if all samples list length is consistent
|
||
seg_len = len(self.ds[0][field])
|
||
for i, ex in enumerate(self.ds):
|
||
val = ex[field]
|
||
assert isinstance(val, list), f"The field '{field}' in the {i}th line is not a list"
|
||
assert len(val) == seg_len, f"The list length in the {i}th line is not consistent"
|
||
|
||
self.field = field
|
||
|
||
def __len__(self):
|
||
return len(self.ds)
|
||
|
||
def __getitem__(self, idx: int):
|
||
return {
|
||
"idx": idx,
|
||
"prompts_list": self.ds[idx][self.field], # List[str]
|
||
}
|
||
|
||
|
||
def cycle(dl):
|
||
while True:
|
||
yield from dl
|