data_loader/video_this_that_dataset.py

import os, sys
import json
import cv2
import math
import shutil
import numpy as np
import random
from PIL import Image
import torch.nn.functional as F
import torch
import os.path as osp
import time
from moviepy.editor import VideoFileClip
from torch.utils.data import Dataset    

# Import files from the local folder
root_path = os.path.abspath('.')
sys.path.append(root_path)
from utils.img_utils import resize_with_antialiasing, numpy_to_pt
from utils.optical_flow_utils import flow_to_image, filter_uv, bivariate_Gaussian
from data_loader.video_dataset import tokenize_captions


# For the 2D dilation
blur_kernel = bivariate_Gaussian(99, 10, 10, 0, grid = None, isotropic = True)


def get_thisthat_sam(config, intput_dir, store_dir = None, flip = False, verbose=False):
    '''
    Args:
        idx (int): The index to the folder we need to process
    '''

    # Read file
    file_path = os.path.join(intput_dir, "data.txt")
    file1 = open(file_path, 'r')
    Lines = file1.readlines()


    # Initial the optical flow format we want
    thisthat_condition = np.zeros((config["video_seq_length"], config["conditioning_channels"], config["height"], config["width"]), dtype=np.float32)  # The last image should be empty


    # Init the image
    sample_img = cv2.imread(os.path.join(intput_dir, "im_0.jpg"))
    org_height, org_width, _ = sample_img.shape

    # Prepare masking
    controlnet_image_index = []
    coordinate_values = []

    # Iterate all points in the txt file
    for idx in range(len(Lines)):

        # Read points
        frame_idx, horizontal, vertical = Lines[idx].split(' ')
        frame_idx, vertical, horizontal = int(frame_idx), int(float(vertical)), int(float(horizontal))

        # Read the mask frame idx
        controlnet_image_index.append(frame_idx)
        coordinate_values.append((vertical, horizontal))


        # Init the base image
        base_img = np.zeros((org_height, org_width, 3)).astype(np.float32)      # Use the original image size
        base_img.fill(255)

        # Draw square around the target position
        dot_range = 10       # Diameter
        for i in range(-1*dot_range, dot_range+1):
            for j in range(-1*dot_range, dot_range+1):
                dil_vertical, dil_horizontal = vertical + i, horizontal + j
                if (0 <= dil_vertical and dil_vertical < base_img.shape[0]) and (0 <= dil_horizontal and dil_horizontal < base_img.shape[1]):
                    if idx == 0:
                        base_img[dil_vertical][dil_horizontal] = [0, 0, 255]        # The first point should be red
                    else:
                        base_img[dil_vertical][dil_horizontal] = [0, 255, 0]        # The second point should be green to distinguish the first point
        
        # Dilate
        if config["dilate"]:
            base_img = cv2.filter2D(base_img, -1, blur_kernel)


        ##############################################################################################################################
        ### The core pipeline of processing is: Dilate -> Resize -> Range Shift -> Transpose Shape -> Store

        # Resize frames  Don't use negative and don't resize in [0,1]
        base_img = cv2.resize(base_img, (config["width"], config["height"]), interpolation = cv2.INTER_CUBIC)


        # Flip the image for aug if needed
        if flip:
            base_img = np.fliplr(base_img)


        # Channel Transform and Range Shift
        if config["conditioning_channels"] == 3:
            # Map to [0, 1] range 
            if store_dir is not None and verbose:    # For the first frame condition visualization
                cv2.imwrite(os.path.join(store_dir, "condition_TT"+str(idx)+".png"), base_img)
            base_img = base_img / 255.0         

        else:
            raise NotImplementedError()

        
        # ReOrganize shape
        base_img = base_img.transpose(2, 0, 1)  # hwc -> chw


        # Check the min max value range
        # if verbose:
        #     print("{} min, max range value is {} - {}".format(intput_dir, np.min(base_img), np.max(base_img)))


        # Write base img based on frame_idx
        thisthat_condition[frame_idx] = base_img        # Only the first frame, the rest is 0 initialized

    ##############################################################################################################################


    if config["motion_bucket_id"] is None:
        # take the motion to stats collected before
        reflected_motion_bucket_id = 200
    else:
        reflected_motion_bucket_id = config["motion_bucket_id"]
    

    # print("Motion Bucket ID is ", reflected_motion_bucket_id)
    return (thisthat_condition, reflected_motion_bucket_id, controlnet_image_index, coordinate_values)    



class Video_ThisThat_Dataset(Dataset):
    '''
        Video Dataset to load sequential frames for training with needed pre-processing and process with optical flow
    '''
    
    def __init__(self, config, device, normalize=True, tokenizer=None):
        # Attribute variables
        self.config = config
        self.device = device
        self.normalize = normalize
        self.tokenizer = tokenizer

        # Obtain values
        self.video_seq_length = config["video_seq_length"]
        self.height = config["height"]
        self.width = config["width"]

        # Process data
        self.video_lists = []
        for dataset_path in config["dataset_path"]:
            for video_name in sorted(os.listdir(dataset_path)):
                if not os.path.exists(os.path.join(dataset_path, video_name, "data.txt")):
                    continue

                self.video_lists.append(os.path.join(dataset_path, video_name))
        print("length of the dataset is ", len(self.video_lists))



        
    def __len__(self):
        return len(self.video_lists)
    

    def _extract_frame_bridge(self, idx, flip=False):
        ''' Extract the frame in video based on the needed fps from already extracted frame
        Args:
            idx (int):                  The index to the file in the directory
            flip (bool):                Bool for whether we will flip
        Returns:
            video_frames (numpy):       Extracted video frames in numpy format
        '''

        # Init the the Video Reader
        # The naming of the Bridge dataset follow a pattern: im_x.jpg, so we need to 
        video_frame_path = self.video_lists[idx]


        # Find needed file
        needed_img_path = []
        for idx in range(self.video_seq_length):
            img_path = os.path.join(video_frame_path, "im_" + str(idx) + ".jpg")
            needed_img_path.append(img_path)



        # Read all img_path based on the order 
        video_frames = []
        for img_path in needed_img_path:
            if not os.path.exists(img_path):
                print("We don't have ", img_path)
            frame = cv2.imread(img_path)

            try:
                frame = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)
            except Exception:
                print("The exception place is ", img_path)
            # Resize frames
            frame = cv2.resize(frame, (self.width, self.height), interpolation = cv2.INTER_CUBIC)

            # Flip aug
            if flip:
                frame = np.fliplr(frame)

            # Collect frames
            video_frames.append(np.expand_dims(frame, axis=0))       # The frame is already RGB, there is no need to convert here.

        
        # Concatenate
        video_frames = np.concatenate(video_frames, axis=0)
        assert(len(video_frames) == self.video_seq_length)

        # Returns
        return video_frames




    def __getitem__(self, idx):
        ''' Get item by idx and pre-process by Resize and Normalize to [0, 1]
        Args:
            idx (int):                  The index to the file in the directory
        Returns:
            return_dict (dict):         video_frames (torch.float32) [-1, 1] and controlnet_condition (torch.float32) [0, 1]
        '''

        # Prepare the text if needed:
        if self.config["use_text"]:
            # Read the file
            file_path = os.path.join(self.video_lists[idx], "lang.txt")
            file = open(file_path, 'r')
            prompt = file.readlines()[0]  # Only read the first line

            if self.config["mix_ambiguous"] and os.path.exists(os.path.join(self.video_lists[idx], "processed_text.txt")):
                # If we don't have this txt file, we skip

                ######################################################## Mix up prompt ########################################################
            
                # Read the file
                file_path = os.path.join(self.video_lists[idx], "processed_text.txt")
                file = open(file_path, 'r')
                prompts = [line for line in file.readlines()]  # Only read the first line

                # Get the componenet
                action = prompts[0][:-1]    
                this = prompts[1][:-1]
                there = prompts[2][:-1]


                random_value = random.random()
                # If less than 0.4, we don't care, just use the most concrete one
                if random_value >= 0.4 and random_value < 0.6:
                    # Mask pick object to "This"
                    prompt = action + " this to " + there
                elif random_value >= 0.6 and random_value < 0.8:
                    # Mask place position to "There"
                    prompt = action + " " + this + " to there"
                elif random_value >= 0.8 and random_value < 1.0:
                    # Just be like "this to there"
                    prompt = action + " this to there"
                
                # print("New prompt is ", prompt)
                ###################################################################################################################################################
            
            # else:
            #     print("We don't have llama processed prompt at ", self.video_lists[idx])
                
        else:
            prompt = ""

        # Tokenize text prompt
        tokenized_prompt = tokenize_captions(prompt, self.tokenizer, self.config)



        # Dataset aug by chance (it is needed to check whether there is any object position words [left|right] in the prompt text)
        flip = False
        if random.random() < self.config["flip_aug_prob"]:
            if self.config["use_text"]:
                if prompt.find("left") == -1 and prompt.find("right") == -1:    # Cannot have position word, like left and right (up and down is ok)
                    flip = True
            else:
                flip = True


        
        # Read frames for different dataset; Currently, we have WebVid / Bridge
        if self.config["dataset_name"] == "Bridge":
            video_frames_raw = self._extract_frame_bridge(idx, flip=flip)
        else:
            raise NotImplementedError("We don't support this dataset loader")


        # Scale [0, 255] -> [-1, 1] if needed
        if self.normalize:
            video_frames = video_frames_raw.astype(np.float32) / 127.5 - 1      # Be careful to cast to float32

        # Transform to Pytorch Tensor in the range [-1, 1]
        video_frames = numpy_to_pt(video_frames)


        # Generate the pairs we need
        intput_dir = self.video_lists[idx]

        # Get the This That point information
        controlnet_condition, reflected_motion_bucket_id, controlnet_image_index, coordinate_values = get_thisthat_sam(self.config, intput_dir, flip=flip)
        controlnet_condition = torch.from_numpy(controlnet_condition)

        # Cast other value to tensor
        reflected_motion_bucket_id = torch.tensor(reflected_motion_bucket_id, dtype=torch.float32) 
        controlnet_image_index = torch.tensor(controlnet_image_index, dtype=torch.int32) 
        coordinate_values = torch.tensor(coordinate_values, dtype=torch.int32)
        

        # The tensor we returned is torch float32. We won't cast here for mixed precision training!
        return {"video_frames" : video_frames, 
                "controlnet_condition" : controlnet_condition, 
                "reflected_motion_bucket_id" : reflected_motion_bucket_id,
                "controlnet_image_index": controlnet_image_index,
                "prompt": tokenized_prompt,
                "coordinate_values": coordinate_values,     # Useless now, but I still passed back
                }