Delete preprocess.py

preprocess.py

@@ -1,260 +0,0 @@
-import torch
-import torchvision
-import json
-import os
-import random 
-import numpy as np
-import argparse
-import decord
-
-from einops import rearrange
-from torchvision import transforms
-from tqdm import tqdm
-from PIL import Image
-from decord import VideoReader, cpu
-from transformers import Blip2Processor, Blip2ForConditionalGeneration
-
-decord.bridge.set_bridge('torch')
-
-class PreProcessVideos:
-    def __init__(
-            self, 
-            config_name,
-            config_save_name,
-            video_directory,
-            random_start_frame,
-            clip_frame_data,
-            max_frames,
-            beam_amount,
-            prompt_amount,
-            min_prompt_length,
-            max_prompt_length,
-            save_dir
-        ):
-
-        # Paramaters for parsing videos
-        self.prompt_amount = prompt_amount
-        self.video_directory = video_directory
-        self.random_start_frame = random_start_frame
-        self.clip_frame_data = clip_frame_data
-        self.max_frames = max_frames
-        self.vid_types = (".mp4", ".avi", ".mov", ".webm", ".flv", ".mjpeg")
-
-        # Parameters for BLIP2
-        self.processor = None
-        self.blip_model = None
-        self.beam_amount = beam_amount
-        self.min_length = min_prompt_length
-        self.max_length = max_prompt_length
-
-        # Helper parameters
-        self.device = "cuda" if torch.cuda.is_available() else "cpu"
-        self.save_dir = save_dir
-
-        # Config parameters
-        self.config_name = config_name
-        self.config_save_name = config_save_name
-
-    # Base dict to hold all the data.
-    # {base_config}
-    def build_base_config(self):
-        return {
-            "name": self.config_name,
-            "data": []
-        }
-
-    # Video dict for individual videos.
-    # {base_config: data -> [{video_path, num_frames, data}]}
-    def build_video_config(self, video_path: str, num_frames: int):
-        return {
-            "video_path": video_path,
-            "num_frames": num_frames,
-            "data": []
-        }
-
-    # Dict for video frames and prompts / captions.
-    # Gets the frame index, then gets a caption for the that frame and stores it.
-    # {base_config: data -> [{name, num_frames, data: {frame_index, prompt}}]}
-    def build_video_data(self, frame_index: int, prompt: str):
-        return {
-            "frame_index": frame_index,
-            "prompt": prompt
-        }
-
-    # Load BLIP2 for processing
-    def load_blip(self):
-        print("Loading BLIP2")
-
-        processor = Blip2Processor.from_pretrained("Salesforce/blip2-opt-2.7b")
-        model = Blip2ForConditionalGeneration.from_pretrained(
-            "Salesforce/blip2-opt-2.7b", torch_dtype=torch.float16
-        )
-        model.to(self.device)
-
-        self.processor = processor
-        self.blip_model = model
-
-    # Process the frames to get the length and image.
-    # The limit parameter ensures we don't get near the max frame length.
-    def video_processor(
-            self, 
-            video_reader: VideoReader, 
-            num_frames: int, 
-            random_start_frame=True,
-            frame_num=0
-        ):
-
-        frame_number = (
-            random.randrange(0, int(num_frames)) if random_start_frame else frame_num
-            )
-        frame = video_reader[frame_number].permute(2,0,1)
-        image = transforms.ToPILImage()(frame).convert("RGB")
-        return frame_number, image
-
-    def get_frame_range(self, derterministic):
-        return range(self.prompt_amount) if self.random_start_frame else derterministic
-
-    def process_blip(self, image: Image):
-        inputs = self.processor(images=image, return_tensors="pt").to(self.device, torch.float16)
-        generated_ids = self.blip_model.generate(
-                **inputs, 
-                num_beams=self.beam_amount, 
-                min_length=self.min_length, 
-                max_length=self.max_length
-            )
-        generated_text = self.processor.batch_decode(
-            generated_ids, 
-            skip_special_tokens=True)[0].strip()
-        
-        return generated_text
-    
-    def get_out_paths(self, prompt, frame_number):
-        out_name= f"{prompt}_{str(frame_number)}"
-        save_path = f"{self.save_dir}/{self.config_save_name}"
-        save_filepath = f"{save_path}/{out_name}.mp4"
-
-        return out_name, save_path, save_filepath
-
-    def save_train_config(self, config: dict):
-        os.makedirs(self.save_dir, exist_ok=True)
-
-        save_json = json.dumps(config, indent=4)
-        save_dir = f"{self.save_dir}/{self.config_save_name}"
-        
-        with open(f"{save_dir}.json", 'w') as f:
-            f.write(save_json)
-
-    def save_video(self, save_path, save_filepath, frames):
-        os.makedirs(save_path, exist_ok=True)
-        torchvision.io.write_video(save_filepath, frames, fps=30)
-
-    # Main loop for processing all videos.
-    def process_videos(self):
-        self.load_blip()
-        config = self.build_base_config()
-
-        if not os.path.exists(self.video_directory):
-            raise ValueError(f"{self.video_directory} does not exist.")
-
-        for _, _, files in tqdm(
-                os.walk(self.video_directory), 
-                desc=f"Processing videos in {self.video_directory}"
-            ):
-            for video in files:
-                if video.endswith(self.vid_types):
-                    video_path = f"{self.video_directory}/{video}"
-                    video_reader = None
-                    derterministic_range = None
-                    video_len = 0
-                    try:
-                        video_reader = VideoReader(video_path, ctx=cpu(0))
-                        video_len = len(video_reader)
-                        frame_step = abs(video_len // self.prompt_amount)
-                        derterministic_range = range(1, abs(video_len - 1), frame_step)          
-                    except:
-                        print(f"Error loading {video_path}. Video may be unsupported or corrupt.")
-                        continue
-                    
-                    # Another try catch block because decord isn't perfect.
-                    try:
-                        num_frames = int(len(video_reader))
-                        video_config = self.build_video_config(video_path, num_frames)
-
-                        # Secondary loop that process a specified amount of prompts, selects a random frame, then appends it.
-                        for i in tqdm(
-                                self.get_frame_range(derterministic_range), 
-                                desc=f"Processing {os.path.basename(video_path)}"
-                            ):
-                            frame_number, image = self.video_processor(
-                                video_reader, 
-                                num_frames, 
-                                self.random_start_frame,
-                                frame_num=i
-                            )
-
-                            prompt = self.process_blip(image)
-                            video_data = self.build_video_data(frame_number, prompt)
-
-                            if self.clip_frame_data:
-
-                                # Minimum value, frame number, max value (length of entire video)
-                                max_range = abs(len(video_reader) - 1)
-                                frame_number = i
-                                frame_number = sorted((1, frame_number, max_range))[1]
-
-                                frame_range = range(frame_number, max_range)
-                                frame_range_nums= list(frame_range)
-
-                                frames = video_reader.get_batch(frame_range_nums[:self.max_frames])
-
-                                out_name, save_path, save_filepath = self.get_out_paths(prompt, frame_number)
-                                
-                                self.save_video(save_path, save_filepath, frames)
-
-                                video_data['clip_path'] = save_filepath
-                                video_config["data"].append(video_data)
-
-                            else:
-                                video_config["data"].append(video_data)
-
-                        config['data'].append(video_config)
-
-                    except Exception as e:
-                        print(e)
-                        continue
-                else:
-                    continue
-
-        print(f"Done. Saving train config to {self.save_dir}.")
-        self.save_train_config(config)
-
-if __name__ == "__main__":
-    parser = argparse.ArgumentParser()
-
-    parser.add_argument('--config_name', help="The name of the configuration.", type=str, default='My Config')
-    parser.add_argument('--config_save_name', help="The name of the config file that's saved.", type=str, default='my_config')
-    parser.add_argument('--video_directory', help="The directory where your videos are located.", type=str, default='./videos')
-    parser.add_argument(
-        '--random_start_frame', 
-        help="Use random start frame when processing videos. Good for long videos where frames have different scenes and meanings.", 
-        action='store_true', 
-        default=False
-    )
-    parser.add_argument(
-        '--clip_frame_data', 
-        help="Save the frames as video clips to HDD/SDD. Videos clips are saved in the same folder as your json directory.", 
-        action='store_true', 
-        default=False
-    )
-    parser.add_argument('--max_frames', help="Maximum frames for clips when --clip_frame_data is enabled.", type=int, default=60)
-    parser.add_argument('--beam_amount', help="Amount for BLIP beam search.", type=int, default=7)
-    parser.add_argument('--prompt_amount', help="The amount of prompts per video that is processed.", type=int, default=25)
-    parser.add_argument('--min_prompt_length', help="Minimum words required in prompt.", type=int, default=15)
-    parser.add_argument('--max_prompt_length', help="Maximum words required in prompt.", type=int, default=30)
-    parser.add_argument('--save_dir', help="The directory to save the config to.", type=str, default=f"{os.getcwd()}/train_data")
-
-    args = parser.parse_args()
-
-    
-    processor = PreProcessVideos(**vars(args))
-    processor.process_videos()