81 lines
2.7 KiB
Python
81 lines
2.7 KiB
Python
from pathlib import Path
|
|
from typing import Dict, List, Optional
|
|
|
|
import numpy as np
|
|
import pytorch_lightning as pl
|
|
import torch
|
|
from pytorch_lightning.utilities.types import EVAL_DATALOADERS
|
|
|
|
from .video_dataset import Annotations
|
|
import json
|
|
|
|
|
|
class ConcatDataset(torch.utils.data.Dataset):
|
|
def __init__(self, datasets):
|
|
self.datasets = datasets
|
|
self.model_id = datasets["reconstruction_dataset"].model_id
|
|
|
|
def __getitem__(self, idx):
|
|
sample = {ds: self.datasets[ds].__getitem__(
|
|
idx) for ds in self.datasets}
|
|
return sample
|
|
|
|
def __len__(self):
|
|
return min(len(self.datasets[d]) for d in self.datasets)
|
|
|
|
|
|
class CustomPromptsDataset(torch.utils.data.Dataset):
|
|
|
|
def __init__(self, prompt_cfg: Dict[str, str]):
|
|
super().__init__()
|
|
|
|
if prompt_cfg["type"] == "prompt":
|
|
self.prompts = [prompt_cfg["content"]]
|
|
elif prompt_cfg["type"] == "file":
|
|
file = Path(prompt_cfg["content"])
|
|
if file.suffix == ".npy":
|
|
self.prompts = np.load(file.as_posix())
|
|
elif file.suffix == ".txt":
|
|
with open(prompt_cfg["content"]) as f:
|
|
lines = [line.rstrip() for line in f]
|
|
self.prompts = lines
|
|
elif file.suffix == ".json":
|
|
with open(prompt_cfg["content"],"r") as file:
|
|
metadata = json.load(file)
|
|
if "videos_root" in prompt_cfg:
|
|
videos_root = Path(prompt_cfg["videos_root"])
|
|
video_path = [str(videos_root / sample["page_dir"] /
|
|
f"{sample['videoid']}.mp4") for sample in metadata]
|
|
else:
|
|
video_path = [str(sample["page_dir"] /
|
|
f"{sample['videoid']}.mp4") for sample in metadata]
|
|
self.prompts = [sample["prompt"] for sample in metadata]
|
|
self.video_path = video_path
|
|
|
|
|
|
|
|
|
|
transformed_prompts = []
|
|
for prompt in self.prompts:
|
|
transformed_prompts.append(
|
|
Annotations.clean_prompt(prompt))
|
|
self.prompts = transformed_prompts
|
|
|
|
def __len__(self):
|
|
return len(self.prompts)
|
|
|
|
def __getitem__(self, index):
|
|
output = {"prompt": self.prompts[index]}
|
|
if hasattr(self,"video_path"):
|
|
output["video"] = self.video_path[index]
|
|
return output
|
|
|
|
|
|
class PromptReader(pl.LightningDataModule):
|
|
def __init__(self, prompt_cfg: Dict[str, str]):
|
|
super().__init__()
|
|
self.predict_dataset = CustomPromptsDataset(prompt_cfg)
|
|
|
|
def predict_dataloader(self) -> EVAL_DATALOADERS:
|
|
return torch.utils.data.DataLoader(self.predict_dataset, batch_size=1, pin_memory=False, shuffle=False, drop_last=False)
|