experiments/process_data.py

import random
import os
from pathlib import Path
import torch
import pandas as pd
import wandb
import time
from tqdm import trange
from torch.utils.data import IterableDataset
from datasets.dummy import DummyVideoDataset
from datasets.openx_base import OpenXVideoDataset
from datasets.droid import DroidVideoDataset
from datasets.something_something import SomethingSomethingDataset
from datasets.epic_kitchen import EpicKitchenDataset
from datasets.pandas import PandasVideoDataset
from datasets.ego4d import Ego4DVideoDataset
from datasets.mixture import MixtureDataset
from datasets.agibot_world import AgibotWorldDataset
from .exp_base import BaseExperiment
from utils.gemini_utils import GeminiCaptionProcessor


class ProcessDataExperiment(BaseExperiment):
    """
    An experiment class for you to easily process an existing
    dataset into another, by creating a new csv metadata file and new files.

    e.g. The `cache_prompt_embed` method illustrates caching the prompt embeddings and
    adding a field `prompt_embed_path` to a copy ofthe metadata csv.

    e.g. The `visualize_dataset` method illustrates visualizing a sample of videos from the dataset with their captions.

    Add your processing methods here, and follow README.md to run.
    """

    compatible_datasets = dict(
        mixture=MixtureDataset,
        mixture_robot=MixtureDataset,
        dummy=DummyVideoDataset,
        something_something=SomethingSomethingDataset,
        epic_kitchen=EpicKitchenDataset,
        pandas=PandasVideoDataset,
        ego4d=Ego4DVideoDataset,
        bridge=OpenXVideoDataset,
        droid=DroidVideoDataset,
        agibot_world=AgibotWorldDataset,
        language_table=OpenXVideoDataset,
        # austin_buds=OpenXVideoDataset,
        # austin_sailor=OpenXVideoDataset,
        # austin_sirius=OpenXVideoDataset,
        # bc_z=OpenXVideoDataset,
        # berkeley_autolab=OpenXVideoDataset,
        # berkeley_cable=OpenXVideoDataset,
        # berkeley_fanuc=OpenXVideoDataset,
        # cmu_stretch=OpenXVideoDataset,
        # dlr_edan=OpenXVideoDataset,
        # dobbe=OpenXVideoDataset,
        # fmb=OpenXVideoDataset,
        # fractal=OpenXVideoDataset,
        # iamlab_cmu=OpenXVideoDataset,
        # jaco_play=OpenXVideoDataset,
        # nyu_franka=OpenXVideoDataset,
        # roboturk=OpenXVideoDataset,
        # stanford_hydra=OpenXVideoDataset,
        # taco_play=OpenXVideoDataset,
        # toto=OpenXVideoDataset,
        # ucsd_kitchen=OpenXVideoDataset,
        # utaustin_mutex=OpenXVideoDataset,
        # viola=OpenXVideoDataset,
    )

    def _build_dataset(
        self, disable_filtering: bool = True, split: str = "all"
    ) -> torch.utils.data.Dataset:
        if disable_filtering:
            self.root_cfg.dataset.filtering.disable = True
        return self.compatible_datasets[self.root_cfg.dataset._name](
            self.root_cfg.dataset, split=split
        )

    def _get_save_dir(self, dataset: torch.utils.data.Dataset):
        save_dir = self.cfg.new_data_root
        if self.cfg.new_data_root is None:
            save_dir = self.output_dir / dataset.data_root.name
        else:
            save_dir = Path(save_dir)
        save_dir.mkdir(parents=True, exist_ok=True)
        return save_dir

    def benchmark_dataloader(self):
        """Benchmark the speed of the dataloader."""
        cfg = self.cfg.benchmark_dataloader
        dataset = self._build_dataset()
        dataloader = torch.utils.data.DataLoader(
            dataset,
            batch_size=cfg.batch_size,
            num_workers=cfg.num_workers,
            shuffle=False,
        )
        for i in trange(1000000):
            time.sleep(0.001)

    def visualize_dataset(self):
        """Visualize a sample of videos from the dataset with their captions.

        This method:
        1. Creates a dataloader for the dataset
        2. Logs the videos and their captions to wandb

        Sample command:
        python main.py +name=process_data experiment=process_data dataset=video_openx experiment.tasks=[visualize_dataset]
        """

        cfg = self.cfg.visualize_dataset
        dataset = self._build_dataset(
            disable_filtering=cfg.disable_filtering, split="training"
        )
        shuffle = not isinstance(dataset, IterableDataset)
        dataloader = torch.utils.data.DataLoader(
            dataset, batch_size=1, num_workers=0, shuffle=shuffle
        )

        log_dict = {}
        self._build_logger()

        samples_seen = 0
        for batch in dataloader:
            if samples_seen >= cfg.n_samples:
                break

            for i in range(len(batch["videos"])):
                if samples_seen >= cfg.n_samples:
                    break

                prompts = None
                if "prompts" in batch:
                    prompts = batch["prompts"][i]

                if cfg.use_processed:
                    video = batch["videos"][i]  # [T, C, H, W]
                    # Convert from [-1, 1] to [0, 255] and correct format for wandb
                    video = ((video + 1) / 2 * 255).clamp(0, 255)
                    video = video.to(torch.uint8).numpy()  # [T, H, W, C]
                    log_dict[f"sample_{samples_seen}"] = wandb.Video(
                        video, caption=prompts, fps=16
                    )
                else:
                    # Log raw video file
                    video_path = str(dataset.data_root / batch["video_path"][i])
                    log_dict[f"sample_{samples_seen}"] = wandb.Video(
                        video_path, caption=prompts, fps=16
                    )

                samples_seen += 1
                if samples_seen % 8 == 0:
                    wandb.log(log_dict)
                    log_dict = {}

        # Log any remaining samples
        if log_dict:
            wandb.log(log_dict)

    def cache_prompt_embed(self):
        """Cache prompt embeddings for all captions in the dataset.

        This method:
        1. Takes captions from the dataset metadata
        2. Generates T5 embeddings for each caption using CogVideo's T5 encoder
        3. Saves embeddings as .pt files alongside the videos
        4. Creates a new metadata CSV with an added 'prompt_embed_path' column

        Sample commands:
        # Cache embeddings for OpenX dataset:
        python main.py +name=process_data experiment=process_data dataset=video_openx experiment.tasks=[cache_prompt_embed]

        # Specify custom output directory:
        python main.py +name=process_data experiment=process_data dataset=video_openx experiment.tasks=[cache_prompt_embed] experiment.new_data_root=data/processed

        # Adjust batch size:
        python main.py +name=process_data experiment=process_data dataset=video_openx experiment.tasks=[cache_prompt_embed] experiment.cache_prompt_embed.batch_size=64
        """
        cfg = self.cfg.cache_prompt_embed
        batch_size = cfg.batch_size

        if self.cfg.num_nodes != 1:
            raise ValueError("This script only supports 1 node. ")

        from algorithms.cogvideo.t5 import T5Encoder

        t5_encoder = T5Encoder(self.root_cfg.algorithm).cuda()
        dataset = self._build_dataset()
        records = dataset.records

        save_dir = self._get_save_dir(dataset)
        metadata_path = save_dir / dataset.metadata_path
        metadata_path.parent.mkdir(parents=True, exist_ok=True)
        print("Saving prompt embeddings and new metadata to ", save_dir)

        new_records = []
        for i in trange(0, len(records), batch_size):
            batch = records[i : i + batch_size]
            prompts = [dataset.id_token + r["caption"] for r in batch]
            embeds = t5_encoder.predict(prompts).cpu()
            for r, embed in zip(batch, embeds):
                video_path = Path(r["video_path"])
                prompt_embed_path = (
                    save_dir / "prompt_embeds" / video_path.with_suffix(".pt")
                )
                prompt_embed_path.parent.mkdir(parents=True, exist_ok=True)
                torch.save(embed.clone(), prompt_embed_path)
                r["prompt_embed_path"] = str(prompt_embed_path.relative_to(save_dir))
                new_records.append(r)

        df = pd.DataFrame.from_records(new_records)
        df.to_csv(metadata_path, index=False)

        print("To review the prompt embeddings, go to ", save_dir)
        print(
            "If everything looks good, you can merge the new dataset into the old "
            "one with the following command:"
        )
        print(f"rsync -av {save_dir}/* {dataset.data_root} && rm -rf {save_dir}")

    def create_gemini_caption(self):
        """
        Create Gemini caption for each video in the dataset.

        1. Init the Dataset, and load all raw records.
        2. Init the GeminiCaptionProcessor with two params:  output_file and num_workers.
        3. Start the processor, and process each record. It will write to the output file.

        For each record in the dataset, it must has "video_path" as the absolute path.
        If each record has some additional keys, like: duration, fps, height, width, n_frames, youtube_key_segment, etc.
        they will be added to the output file. Check "Class VideoEntry" below for more details.

        Sample command:
        python main.py +name=create_gemini_caption experiment=process_data dataset=pandas experiment.tasks=[create_gemini_caption]
        """
        cfg = self.cfg.create_gemini_caption
        num_workers = cfg.n_workers

        dataset = self._build_dataset()
        records = dataset.records

        save_dir = self._get_save_dir(dataset)
        metadata_path = dataset.metadata_path.with_suffix(".json")
        metadata_path = metadata_path.parent / ("gemini_" + metadata_path.name)
        output_file = save_dir / metadata_path

        for r in records:
            r["video_path"] = str((dataset.data_root / r["video_path"]).absolute())

        if not os.path.exists(records[0]["video_path"]):
            raise ValueError("video_path must be an absolute path")

        processor = GeminiCaptionProcessor(output_file, num_workers=num_workers)
        processor.process_entries(records)
        print("To review the captions, go to ", output_file)
        print(
            "If everything looks good, you can merge the new dataset into the old "
            "one with the following command:"
        )
        print(f"rsync -av {save_dir}/* {dataset.data_root} && rm -rf {save_dir}")

    def run_hand_pose_estimation(self):

        import queue
        import threading
        import decord

        # see https://github.com/ibaiGorordo/Sapiens-Pytorch-Inference/blob/main/image_pose_estimation.py
        from sapiens_inference import SapiensPoseEstimation, SapiensPoseEstimationType
        import time

        # also use confidence score > 0.3
        # for each key, it will store x, y, confidence score
        hand_keypoints_keys_list = [
            # in total of 40 keypoints
            # Right hand
            "right_wrist",
            "right_thumb4",
            "right_thumb3",
            "right_thumb2",
            "right_thumb_third_joint",
            "right_forefinger4",
            "right_forefinger3",
            "right_forefinger2",
            "right_forefinger_third_joint",
            "right_middle_finger4",
            "right_middle_finger3",
            "right_middle_finger2",
            "right_middle_finger_third_joint",
            "right_ring_finger4",
            "right_ring_finger3",
            "right_ring_finger2",
            "right_ring_finger_third_joint",
            "right_pinky_finger4",
            "right_pinky_finger3",
            "right_pinky_finger2",
            "right_pinky_finger_third_joint",
            # Left hand
            "left_wrist",
            "left_thumb4",
            "left_thumb3",
            "left_thumb2",
            "left_thumb_third_joint",
            "left_forefinger4",
            "left_forefinger3",
            "left_forefinger2",
            "left_forefinger_third_joint",
            "left_middle_finger4",
            "left_middle_finger3",
            "left_middle_finger2",
            "left_middle_finger_third_joint",
            "left_ring_finger4",
            "left_ring_finger3",
            "left_ring_finger2",
            "left_ring_finger_third_joint",
            "left_pinky_finger4",
            "left_pinky_finger3",
            "left_pinky_finger2",
            "left_pinky_finger_third_joint",
        ]

        cfg = self.cfg.run_hand_pose_estimation

        dataset = self._build_dataset()
        records = dataset.records

        # for debug, only process 50 videos
        # records = records[:50]
        # random sample 50 videos
        records = random.sample(records, 50)
        save_dir = self._get_save_dir(dataset)
        Path(save_dir).mkdir(parents=True, exist_ok=True)
        print(f"Saving hand pose estimation results to {save_dir}")

        # Create queues for communication between producer and consumer
        frame_queue = queue.Queue(
            maxsize=100
        )  # Limit queue size to prevent memory issues
        STOP_TOKEN = "DONE"

        def producer(records, data_root):
            for record in records:
                try:
                    video_path = Path(data_root) / record["video_path"]
                    vr = decord.VideoReader(str(video_path))
                    n_frames = len(vr)

                    if n_frames == 0:
                        print(f"No frames found in {record['video_path']}")
                        continue

                    # Get first, middle, and last frame indices
                    frame_indices = [0, n_frames // 2, n_frames - 1]
                    frames = vr.get_batch(
                        frame_indices
                    ).asnumpy()  # Shape: (3, H, W, C)
                    # also resize each frame to 768x1024. with height 768, width 1024
                    # frames = [cv2.resize(frame, (1024, 768)) for frame in frames]
                    # Put frames and relative path in queue
                    frame_queue.put(
                        {
                            "frames": frames,
                            "video_path": str(
                                record["video_path"]
                            ),  # Keep relative path
                            "frame_indices": frame_indices,  # Keep track of which frames
                        }
                    )
                except Exception as e:
                    print(f"Error processing {record['video_path']}: {e}")
                    continue

            # Signal completion
            frame_queue.put(STOP_TOKEN)

        start_time = time.time()
        # Start producer thread
        producer_thread = threading.Thread(
            target=producer, args=(records, dataset.data_root), daemon=True
        )
        producer_thread.start()

        # Initialize the pose estimator
        dtype = torch.float16
        estimator = SapiensPoseEstimation(
            SapiensPoseEstimationType.POSE_ESTIMATION_03B, dtype=dtype
        )

        # Prepare a list to collect results
        # Each result will be a dict with video_path, frame_index, keypoints
        results = []

        while True:
            item = frame_queue.get()
            if item == STOP_TOKEN:
                break

            frames = item["frames"]  # Shape: (3, H, W, C)
            video_path = item["video_path"]
            frame_indices = item.get("frame_indices", [0, 1, 2])

            ret_per_video = {
                "video_path": video_path,
                "frame_indices": frame_indices,
                "keypoints_list": [],
            }
            for idx, frame in zip(frame_indices, frames):
                try:
                    # Convert frame from BGR (OpenCV) to RGB
                    # frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
                    frame_rgb = frame

                    # Run pose estimation
                    result_img, keypoints = estimator(frame_rgb)

                    # Optionally, you can save or display the result_img
                    # For example, to save the annotated image:
                    # annotated_img_path = Path(save_dir) / f"{Path(video_path).stem}_frame_{idx}.jpg"
                    # cv2.imwrite(str(annotated_img_path), cv2.cvtColor(result_img, cv2.COLOR_RGB2BGR))

                    # Flatten keypoints and prepare the result entry
                    # Assuming keypoints is a NumPy array of shape (num_keypoints, 2) or similar
                    # print("debug", keypoints)
                    keypoints_flat = keypoints  # list of dict.

                    # only store the keypoints that are in hand_keypoints_keys_list
                    keypoints_flat = [
                        {
                            k: kp_dict[k]
                            for k in hand_keypoints_keys_list
                            if k in kp_dict
                        }
                        for kp_dict in keypoints_flat
                    ]

                    # then remove pred whose confidence score is less than 0.3
                    keypoints_flat = [
                        {k: v for k, v in kp_dict.items() if v[2] > 0.3}
                        for kp_dict in keypoints_flat
                    ]
                    result_entry = {
                        "frame_index": idx,
                        "keypoints_list": keypoints_flat,
                        "num_keypoints": sum([len(_) for _ in keypoints_flat]),
                    }

                    ret_per_video["keypoints_list"].append(result_entry)

                except Exception as e:
                    print(
                        f"Error running pose estimation for frame {idx} of {video_path}: {e}"
                    )
                    continue

            # tell if there exists any keypoints in the video, if not skip the video
            num_keypoints = sum(
                [_.get("num_keypoints", 0) for _ in ret_per_video["keypoints_list"]]
            )
            if num_keypoints > 0:
                results.append(ret_per_video)
            frame_queue.task_done()

        producer_thread.join()

        end_time = time.time()
        print(f"Time taken: {end_time - start_time} seconds")
        print(f"Total number of videos processed with keypoints: {len(results)}")

        # Convert results to JSON format
        if results:
            # Each result already contains:
            # - video_path
            # - frame_index
            # - keypoints_list (list of dictionaries with pose data)

            # Save to JSON
            json_path = Path(save_dir) / "hand_pose_results.json"
            import json

            with open(json_path, "w") as f:
                json.dump(results, f, indent=2)
            print(f"Results saved to {json_path}")
        else:
            print("No results to save.")

    def run_human_detection(self):

        import queue
        import threading
        import decord
        from utils.detector_utils import Detector
        import time

        detector = Detector()  # bboxes = detector.detech(np_img_BGR)

        cfg = self.cfg.run_human_detection

        dataset = self._build_dataset()
        records = dataset.records
        # try 40k videos for now.
        # records = records[:40000]

        num_workers = cfg.total_workers
        job_id = cfg.job_id

        records = records[job_id::num_workers]
        # for debug, only process 50 videos
        # records = records[:50]
        # random sample 50 videos
        # records = random.sample(records, 50)
        save_dir = self._get_save_dir(dataset)
        Path(save_dir).mkdir(parents=True, exist_ok=True)
        print(f"Saving hand pose estimation results to {save_dir}")

        # Create queues for communication between producer and consumer
        frame_queue = queue.Queue(
            maxsize=100
        )  # Limit queue size to prevent memory issues
        STOP_TOKEN = "DONE"

        def producer(records, data_root):
            for record in records:
                try:
                    video_path = Path(data_root) / record["video_path"]
                    vr = decord.VideoReader(str(video_path))
                    n_frames = len(vr)

                    if n_frames == 0:
                        print(f"No frames found in {record['video_path']}")
                        continue

                    # get one frame every second, read fps first then get frame indices
                    fps = vr.get_avg_fps()
                    frame_indices = [int(i * fps) for i in range(int(n_frames // fps))]
                    frames = vr.get_batch(
                        frame_indices
                    ).asnumpy()  # Shape: (n_f, H, W, C)
                    # also resize each frame to 768x1024. with height 768, width 1024
                    # frames = [cv2.resize(frame, (1024, 768)) for frame in frames]
                    # Put frames and relative path in queue
                    frame_queue.put(
                        {
                            "frames": frames,
                            "video_path": str(
                                record["video_path"]
                            ),  # Keep relative path
                            "frame_indices": frame_indices,  # Keep track of which frames
                        }
                    )
                except Exception as e:
                    print(f"Error processing {record['video_path']}: {e}")
                    continue

            # Signal completion
            frame_queue.put(STOP_TOKEN)

        start_time = time.time()
        # Start producer thread
        producer_thread = threading.Thread(
            target=producer, args=(records, dataset.data_root), daemon=True
        )
        producer_thread.start()

        # Initialize the pose estimator
        dtype = torch.float16

        # Prepare a list to collect results
        # Each result will be a dict with video_path, frame_index, keypoints
        results = []

        while True:
            item = frame_queue.get()
            if item == STOP_TOKEN:
                break

            frames = item["frames"]  # Shape: (3, H, W, C)
            video_path = item["video_path"]
            frame_indices = item.get("frame_indices", [0, 1, 2])

            ret_per_video = {
                "video_path": video_path,
                "frame_indices": frame_indices,
                "bbox_list": [],
            }
            num_detections = 0
            for idx, frame in zip(frame_indices, frames):
                try:
                    bboxes = detector.detect(
                        frame
                    ).tolist()  # [(x1, y1, x2, y2), ...] or empty list []
                    ret_per_video["bbox_list"].append(bboxes)
                    num_detections += len(bboxes)
                except Exception as e:
                    print(
                        f"Error running human detection for frame {idx} of {video_path}: {e}"
                    )
                    continue

            results.append(ret_per_video)
            frame_queue.task_done()

        producer_thread.join()

        end_time = time.time()
        print(f"Time taken: {end_time - start_time} seconds")
        print(f"Total number of videos processed with human detections: {len(results)}")

        # Convert results to JSON format
        if results:
            # Each result already contains:
            # - video_path
            # - frame_index
            # - bbox_list (list of list of bbox)

            # Save to JSON
            json_path = Path(save_dir) / f"human_detection_results_{job_id}.json"
            import json

            with open(json_path, "w") as f:
                json.dump(results, f, indent=2)
            print(f"Results saved to {json_path}")
        else:
            print("No results to save.")