348 lines
14 KiB
Python
348 lines
14 KiB
Python
# -*- coding: utf-8 -*-
|
|
# Copyright (c) Alibaba, Inc. and its affiliates.
|
|
import csv
|
|
import os
|
|
import shutil
|
|
import time
|
|
import decord
|
|
from tqdm import tqdm
|
|
import gradio as gr
|
|
|
|
from scepter.modules.utils.directory import get_md5
|
|
from scepter.modules.utils.file_system import FS
|
|
from scepter.studio.preprocess.caption_editor_ui.component_names import \
|
|
Text2VideoDataCardName
|
|
from scepter.studio.preprocess.utils.data_card import (BaseDataCard, find_prefix)
|
|
|
|
|
|
class Text2VideoDataCard(BaseDataCard):
|
|
def __init__(self,
|
|
dataset_folder,
|
|
dataset_name=None,
|
|
src_file=None,
|
|
surfix=None,
|
|
user_name='admin',
|
|
language='en'
|
|
):
|
|
super().__init__(dataset_folder,
|
|
dataset_name=dataset_name,
|
|
user_name=user_name)
|
|
self.meta['task_type'] = 'txt2vid'
|
|
self.components_name = Text2VideoDataCardName(language)
|
|
if self.new_dataset:
|
|
if surfix == '.zip':
|
|
file_list = self.load_from_zip(src_file, dataset_folder,
|
|
self.local_dataset_folder)
|
|
elif surfix in ['.txt', '.csv']:
|
|
file_list = self.load_from_list(src_file, dataset_folder,
|
|
self.local_dataset_folder)
|
|
elif surfix is None:
|
|
file_list = []
|
|
else:
|
|
raise gr.Error(
|
|
f'{self.components_name.illegal_data_err2} {surfix}')
|
|
|
|
is_flag = FS.put_dir_from_local_dir(self.local_dataset_folder,
|
|
dataset_folder,
|
|
multi_thread=True)
|
|
if not is_flag:
|
|
raise gr.Error(f'{self.components_name.illegal_data_err3}')
|
|
|
|
self.meta['cursor'] = 0 if len(file_list) > 0 else -1
|
|
self.meta['file_list'] = file_list
|
|
self.update_dataset()
|
|
|
|
def apply_changes(self):
|
|
edit_index_list = self.edit_list
|
|
for index in edit_index_list:
|
|
one_data = self.data[index]
|
|
self.data[index]['caption'] = one_data['edit_caption']
|
|
|
|
self.update_dataset()
|
|
return True, ''
|
|
|
|
def load_from_list(self, save_file, dataset_folder, local_dataset_folder):
|
|
file_list = []
|
|
videos_folder = os.path.join(local_dataset_folder, 'videos')
|
|
os.makedirs(videos_folder, exist_ok=True)
|
|
with FS.get_from(save_file) as local_path:
|
|
with open(local_path, 'r') as f:
|
|
for line in tqdm(f):
|
|
line = line.strip()
|
|
if line == '':
|
|
continue
|
|
try:
|
|
src_video_path, caption = line.split(
|
|
'#;#', 1)
|
|
except Exception:
|
|
try:
|
|
src_video_path, caption = line.split(
|
|
',', 1)
|
|
except Exception:
|
|
raise gr.Error(
|
|
self.components_name.illegal_data_err1)
|
|
relative_path = os.path.join(
|
|
'videos', f'{get_md5(src_video_path)[:18]}_{int(time.time())}.mp4')
|
|
video_path = os.path.join(dataset_folder, relative_path)
|
|
FS.get_from(src_video_path, local_path=video_path)
|
|
is_legal, new_path, prefix = find_prefix(src_video_path)
|
|
w, h, fps, duration = self.get_video_meta(video_path)
|
|
|
|
file_list.append({
|
|
'video_path':
|
|
video_path,
|
|
'relative_path':
|
|
relative_path,
|
|
'width':
|
|
w,
|
|
'height':
|
|
h,
|
|
'fps':
|
|
fps,
|
|
'duration':
|
|
duration,
|
|
'caption':
|
|
caption,
|
|
'edit_caption':
|
|
caption,
|
|
'prefix':
|
|
prefix
|
|
})
|
|
return file_list
|
|
|
|
def load_from_zip(self, save_file, data_folder, local_dataset_folder):
|
|
with FS.get_from(save_file) as local_path:
|
|
res = os.popen(
|
|
f"unzip -o '{local_path}' -d '{local_dataset_folder}'")
|
|
res = res.readlines()
|
|
if not os.path.exists(local_dataset_folder):
|
|
raise gr.Error(f'Unzip {save_file} failed {str(res)}')
|
|
file_folder = None
|
|
train_list = None
|
|
hit_dir = None
|
|
raw_list = {}
|
|
mac_osx = os.path.join(local_dataset_folder, '__MACOSX')
|
|
if os.path.exists(mac_osx):
|
|
res = os.popen(f"rm -rf '{mac_osx}'")
|
|
res = res.readlines()
|
|
for one_dir in FS.walk_dir(local_dataset_folder, recurse=False):
|
|
if one_dir.endswith('__MACOSX'):
|
|
res = os.popen(f"rm -rf '{one_dir}'")
|
|
res = res.readlines()
|
|
continue
|
|
if FS.isdir(one_dir):
|
|
if one_dir.endswith('videos') or one_dir.endswith('videos/'):
|
|
file_folder = one_dir
|
|
hit_dir = one_dir
|
|
else:
|
|
sub_dir = FS.walk_dir(one_dir)
|
|
for one_s_dir in sub_dir:
|
|
if FS.isdir(one_s_dir) and one_s_dir.split(
|
|
one_dir)[1].replace('/', '') == 'videos':
|
|
file_folder = one_s_dir
|
|
hit_dir = one_dir
|
|
if FS.isfile(one_s_dir) and one_s_dir.split(
|
|
one_dir)[1].replace('/', '') == 'train.csv':
|
|
train_list = one_s_dir
|
|
if file_folder is not None and train_list is not None:
|
|
break
|
|
elif one_dir.endswith('train.csv'):
|
|
train_list = one_dir
|
|
else:
|
|
continue
|
|
if file_folder is not None and train_list is not None:
|
|
break
|
|
if file_folder is None and len(raw_list) < 1:
|
|
raise gr.Error(
|
|
"video doesn't exist, or nothing exists in your zip")
|
|
|
|
if train_list is None:
|
|
raise gr.Error("pair list doesn't exist")
|
|
new_file_folder = f'{local_dataset_folder}/videos'
|
|
os.makedirs(new_file_folder, exist_ok=True)
|
|
|
|
if file_folder is not None:
|
|
_ = FS.get_dir_to_local_dir(file_folder, new_file_folder)
|
|
|
|
if not os.path.exists(new_file_folder):
|
|
raise gr.Error(f'{str(res)}')
|
|
new_train_list = f'{local_dataset_folder}/train.csv'
|
|
res = os.popen(f"mv '{train_list}' '{new_train_list}'")
|
|
res = res.readlines()
|
|
if not os.path.exists(new_train_list):
|
|
raise gr.Error(f'{str(res)}')
|
|
if not file_folder == hit_dir:
|
|
try:
|
|
res = os.popen(f"rm -rf '{hit_dir}/videos/*'")
|
|
_ = res.readlines()
|
|
res = os.popen(f"rm -rf '{hit_dir}'")
|
|
_ = res.readlines()
|
|
|
|
except Exception:
|
|
pass
|
|
file_list = self.load_train_file(new_train_list)
|
|
# remove unused data
|
|
for one_dir in FS.walk_dir(local_dataset_folder):
|
|
if 'videos' in one_dir or one_dir.endswith(
|
|
'file.csv') or one_dir.endswith('train.csv'):
|
|
continue
|
|
os.system(f'rm -rf {one_dir}')
|
|
return file_list
|
|
|
|
def load_train_file(self, file_path):
|
|
base_folder = os.path.dirname(file_path)
|
|
file_list = []
|
|
video_set = set()
|
|
|
|
with open(file_path, 'r') as f:
|
|
reader = csv.reader(f)
|
|
for row in reader:
|
|
if len(row) == 2:
|
|
src_video_path, prompt = row[0], row[1]
|
|
else:
|
|
return gr.Error(self.components_name.illegal_data_err2)
|
|
if src_video_path == 'Target:FILE':
|
|
continue
|
|
|
|
local_video_path = os.path.join(base_folder, src_video_path)
|
|
w, h, fps, duration = self.get_video_meta(local_video_path)
|
|
if src_video_path in video_set:
|
|
src_video_path, surfix = os.path.splitext(src_video_path)
|
|
src_video_path = f'{src_video_path}_{int(time.time() * 100)}{surfix}'
|
|
new_local_video_path = os.path.join(
|
|
base_folder, src_video_path)
|
|
self.copy_video(src_video_path, new_local_video_path)
|
|
video_set.add(src_video_path)
|
|
|
|
file_list.append({
|
|
'video_path':
|
|
local_video_path,
|
|
'relative_path':
|
|
src_video_path,
|
|
'width':
|
|
w,
|
|
'height':
|
|
h,
|
|
'fps':
|
|
fps,
|
|
"duration":
|
|
duration,
|
|
'caption':
|
|
prompt,
|
|
'edit_caption':
|
|
prompt,
|
|
'prefix':
|
|
''
|
|
})
|
|
return file_list
|
|
|
|
def write_train_file(self):
|
|
file_list = self.meta['file_list']
|
|
with open(self.local_train_file, 'w') as f:
|
|
writer = csv.writer(f)
|
|
writer.writerow(['Target:FILE', 'Prompt'])
|
|
for one_file in file_list:
|
|
relative_file = one_file['relative_path']
|
|
if relative_file.startswith('/'):
|
|
relative_file = relative_file[1:]
|
|
writer.writerow([relative_file, one_file['caption'].
|
|
strip().replace("\n", "")])
|
|
FS.put_object_from_local_file(self.local_train_file, self.train_file)
|
|
|
|
def write_data_file(self):
|
|
file_list = self.meta['file_list']
|
|
with open(self.local_save_file_list, 'w') as f:
|
|
for one_file in file_list:
|
|
f.write('{}#;#{}#;#{}#;#{}\n'.format(one_file['relative_path'],
|
|
one_file['width'],
|
|
one_file['height'],
|
|
one_file['caption'].strip().replace("\n", ""))
|
|
)
|
|
FS.put_object_from_local_file(self.local_save_file_list,
|
|
self.save_file_list)
|
|
|
|
def add_record(self, video, caption, **kwargs):
|
|
local_work_dir = self.meta['local_work_dir']
|
|
work_dir = self.meta['work_dir']
|
|
|
|
save_folder = os.path.join(local_work_dir, 'videos')
|
|
os.makedirs(save_folder, exist_ok=True)
|
|
w, h, fps, duration = self.get_video_meta(video)
|
|
|
|
relative_path = os.path.join(
|
|
'videos', f'{get_md5(video)[:18]}_{int(time.time())}.mp4')
|
|
video_path = os.path.join(work_dir, relative_path)
|
|
local_video_path = os.path.join(local_work_dir, relative_path)
|
|
self.copy_video(video, local_video_path)
|
|
|
|
self.data.append({
|
|
'video_path': video_path,
|
|
'relative_path': relative_path,
|
|
'width': w,
|
|
'height': h,
|
|
'fps': fps,
|
|
'duration': duration,
|
|
'caption': caption,
|
|
'edit_caption': caption,
|
|
'prefix': ''
|
|
})
|
|
|
|
self.set_cursor(len(self.meta['file_list']) - 1)
|
|
self.update_dataset()
|
|
return True
|
|
|
|
def copy_video(self, source_path, target_path):
|
|
if not os.path.isfile(source_path):
|
|
raise gr.Error('Video path not exist.')
|
|
try:
|
|
shutil.copy2(source_path, target_path)
|
|
except Exception as e:
|
|
raise gr.Error(str(e))
|
|
|
|
def get_video_meta(self, video):
|
|
video_reader = decord.VideoReader(video)
|
|
w = video_reader[0].shape[1]
|
|
h = video_reader[0].shape[0]
|
|
fps = video_reader.get_avg_fps()
|
|
video_length = len(video_reader)
|
|
duration = video_length / fps
|
|
|
|
return w, h, fps, duration
|
|
|
|
def delete_record(self):
|
|
if len(self) < 1:
|
|
raise gr.Error(self.components_name.delete_err1)
|
|
current_file = self.data.pop(self.cursor)
|
|
self.set_cursor(self.cursor - 1)
|
|
local_file = os.path.join(self.meta['local_work_dir'],
|
|
current_file['relative_path'])
|
|
try:
|
|
os.remove(local_file)
|
|
except Exception:
|
|
print(f'remove file {local_file} error')
|
|
|
|
if self.cursor >= len(self.meta['file_list']):
|
|
self.set_cursor(0)
|
|
if self.cursor < 0:
|
|
self.set_cursor(len(self) - 1)
|
|
if len(self.meta['file_list']) == 0:
|
|
self.set_cursor(-1)
|
|
self.update_dataset()
|
|
|
|
def export_zip(self, export_folder):
|
|
self.update_dataset()
|
|
zip_path = os.path.join(export_folder, f'{self.dataset_name}.zip')
|
|
local_zip, _ = FS.map_to_local(zip_path)
|
|
os.makedirs(os.path.dirname(local_zip), exist_ok=True)
|
|
res = os.popen(
|
|
f"cd '{self.local_work_dir}' && mkdir -p '{self.dataset_name}' "
|
|
f"&& cp -rf videos '{self.dataset_name}/videos' "
|
|
f"&& cp -rf train.csv '{self.dataset_name}/train.csv' "
|
|
f"&& zip -r '{os.path.abspath(local_zip)}' '{self.dataset_name}'/* "
|
|
f"&& rm -rf '{self.dataset_name}'")
|
|
print(res.readlines())
|
|
FS.put_object_from_local_file(local_zip, zip_path)
|
|
if not FS.exists(zip_path):
|
|
raise gr.Error(self.components_name.export_zip_err1)
|
|
return local_zip
|