# -*- coding: utf-8 -*- # Copyright (c) Alibaba, Inc. and its affiliates. import copy import csv import datetime import json import os.path import gradio as gr import imagehash from PIL import Image from tqdm import tqdm 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 \ Text2ImageDataCardName def find_prefix(file_path): for k in FS._prefix_to_clients.keys(): if file_path.startswith(k): return True, file_path, '' elif FS.exists(os.path.join(k, file_path)): return True, os.path.join(k, file_path), '' elif FS.exists(os.path.join(k, 'datasets', file_path)): return True, os.path.join(k, 'datasets', file_path), 'datasets' return False, None, None def del_prefix(file_path, prefix=''): for k in FS._prefix_to_clients.keys(): if file_path.startswith(k): file_path = file_path.replace(k, '') while file_path.startswith('/'): file_path = file_path[1:] if not prefix == '' and file_path.startswith(prefix): file_path = file_path.split(prefix)[-1] while file_path.startswith('/'): file_path = file_path[1:] return True, file_path return False, file_path def get_image_meta(image): img = Image.open(image) return img.size class BaseDataCard(object): def __init__(self, dataset_folder, dataset_name=None, user_name='admin', batch=8): self.dataset_folder = dataset_folder self.batch = batch self.local_dataset_folder, _ = FS.map_to_local(dataset_folder) self.meta_file = os.path.join(dataset_folder, 'meta.json') self.local_meta = os.path.join(self.local_dataset_folder, 'meta.json') self.train_file = os.path.join(dataset_folder, 'train.csv') self.local_train_file = os.path.join(self.local_dataset_folder, 'train.csv') self.save_file_list = os.path.join(dataset_folder, 'file.txt') self.local_save_file_list = os.path.join(self.local_dataset_folder, 'file.txt') self.edit_list = [] self.edit_cursor = -1 current_time = self.get_time() ''' A legal meta is like as follows: { 'dataset_name': '', 'cursor': cursor, 'file_list': [], # use to manage all data status 'train_csv': train_csv, 'save_file_list': save_file_list, 'is_valid': True or False , 'create_time': 'YYYYMMDD-HHMMSS', 'update_time': 'YYYYMMDD-HHMMSS', 'user_name': '' } ''' self.default_meta = { 'dataset_name': dataset_name, 'cursor': -1, 'file_list': [], # use to manage all data status 'train_csv': self.train_file, 'save_file_list': self.save_file_list, 'is_valid': True, 'create_time': current_time, 'update_time': current_time, 'user_name': user_name } if FS.exists(self.meta_file): FS.get_dir_to_local_dir(self.dataset_folder, self.local_dataset_folder) if FS.exists(self.local_meta): self.meta = self.check_legal_dataset( json.load(open(self.local_meta, 'r'))) else: self.meta = {} else: self.meta = {} if len(self.meta) < 1: assert dataset_name is not None self.meta = self.default_meta self.meta['dataset_name'] = dataset_name os.makedirs(self.local_dataset_folder, exist_ok=True) self.new_dataset = True else: self.new_dataset = False self.meta['local_work_dir'] = self.local_dataset_folder self.meta['work_dir'] = dataset_folder self.dataset_name = self.meta['dataset_name'] self.start_cursor = (self.cursor // batch) * batch def check_legal_dataset(self, meta): # If meta has dataset, this dataset is legal. # otherwise, we take it as a new dataset or illegal one. if 'dataset_name' not in meta: return {} is_update = False for key, value in self.default_meta.items(): if key not in meta: meta[key] = self.default_meta[key] is_update = True if is_update: self.save_meta(meta) return meta def get_time(self): return '{0:%Y%m%d%-H%M%S}'.format(datetime.datetime.now()) def update_dataset(self): self.write_train_file() self.write_data_file() self.save_meta() def save_meta(self, meta=None): meta = self.meta if meta is None else meta save_meta = copy.deepcopy(meta) if 'local_work_dir' in meta: save_meta.pop('local_work_dir') if 'work_dir' in meta: save_meta.pop('work_dir') json.dump(save_meta, open(self.local_meta, 'w')) FS.put_object_from_local_file(self.local_meta, self.meta_file) return True def write_train_file(self): raise NotImplementedError def write_data_file(self): raise NotImplementedError @property def cursor(self): return self.meta['cursor'] def __len__(self): return len(self.meta['file_list']) @property def data(self): return self.meta['file_list'] def set_cursor(self, cursor): if cursor >= len(self): self.meta['cursor'] = 0 elif len(self) == 0: self.meta['cursor'] = -1 else: self.meta['cursor'] = cursor # if self.cursor is the start of current batch or end of current batch # we employ stride forward or backward automaticlly. # but now it’s not compatiable with gradio # if self.cursor - self.start_cursor == self.batch - 1: # self.start_cursor = self.cursor # elif self.cursor - self.start_cursor == 0: # self.start_cursor = self.cursor - self.batch + 1 # if self.start_cursor < 0: self.start_cursor = 0 # else: # self.start_cursor = (self.cursor//self.batch) * self.batch def deactive_dataset(self): self.meta['is_valid'] = False self.update_dataset() def active_dataset(self): self.meta['is_valid'] = True self.update_dataset() @property def is_valid(self): return self.meta['is_valid'] @property def user_name(self): return self.meta['user_name'] @property def work_dir(self): return self.meta['work_dir'] @property def local_work_dir(self): return self.meta['local_work_dir'] @property def get_batch(self): end_index = self.start_cursor + self.batch if end_index >= len(self): extend_data = self.data[0:end_index - len(self) + 1] else: extend_data = [] return self.data[self.start_cursor:end_index] + extend_data def set_edit_range(self, range_list): self.edit_list = [] illegal_tup = [] if range_list == -1: self.edit_list = list(range(len(self))) self.samples_list = self.data return True, '' for range_tup in range_list.split(','): if range_tup.strip() == '': continue if '-' in range_tup: num_tup = range_tup.split('-') if not len(num_tup) == 2: illegal_tup.append( f"{range_tup} is illegal, more than one '-'") continue try: start_num = int(num_tup[0]) end_num = int(num_tup[1]) except Exception as e: illegal_tup.append( f'{range_tup} is illegal, start number ' f'or end number is not int number, error {e}. ') continue if start_num < 1 or end_num < 1: illegal_tup.append(f'{range_tup} is illegal, start number ' f'or end number should >= 1. ') continue if start_num > len(self) or end_num > len(self): illegal_tup.append( f'{range_tup} is illegal, start number ' f'or end number should <= length of this dataset {len(self)}. ' ) continue self.edit_list.extend( list(range(start_num - 1, end_num - 1, 1))) else: try: num = int(range_tup) except Exception as e: illegal_tup.append( f'{range_tup} is illegal, number is not int number,' f' error {e}. ') continue if num < 1: illegal_tup.append(f'{range_tup} is illegal, number ' f'should >= 1. ') continue if num > len(self): illegal_tup.append( f'{range_tup} is illegal, number ' f'should <= length of this dataset {len(self)}. ') continue self.edit_list.append(num - 1) self.edit_list.sort() self.samples_list = [self.data[index] for index in self.edit_list] if len(self.samples_list) < 1: illegal_tup.append('select list is blank.') if len(illegal_tup) > 0: return False, ' '.join(illegal_tup) else: return True, '' @property def edit_samples(self): return self.samples_list def edit_index_from_cursor(self, cursor): if cursor in self.edit_list: return self.edit_list.index(cursor) return -1 def cursor_from_edit_index(self, index): if index < 0 or index >= len(self.edit_list): return -1 return self.edit_list[index] def apply_changes(self): edit_index_list = self.edit_list for index in edit_index_list: one_data = self.data[index] relative_image_path = one_data['relative_path'] local_image_path = os.path.join(self.meta['local_work_dir'], relative_image_path) image_path = one_data['image_path'] edit_relative_image_path = one_data.get('edit_relative_path', one_data['relative_path']) local_edit_image_path = os.path.join(self.meta['local_work_dir'], edit_relative_image_path) if not relative_image_path == edit_relative_image_path: try: os.rename(local_edit_image_path, local_image_path) except Exception as e: msg = f'Apply edited image failed, error is {e}' return False, msg FS.put_object_from_local_file(local_image_path, image_path) try: os.remove(local_edit_image_path) except Exception: pass self.data[index]['edit_relative_path'] = one_data['relative_path'] self.data[index]['edit_image_path'] = one_data['image_path'] self.data[index]['caption'] = one_data['edit_caption'] self.data[index]['width'] = one_data['edit_width'] self.data[index]['height'] = one_data['edit_height'] self.update_dataset() return True, '' # @property # def select_index(self): # return self.cursor - self.start_cursor @property def current_record(self): if self.cursor >= len(self): self.set_cursor(0) if self.cursor >= 0: return self.meta['file_list'][self.cursor] else: return {} def modify_data_name(self, new_dataset_name): self.meta['dataset_name'] = new_dataset_name self.update_dataset() def edit_caption(self, edit_caption): if self.cursor >= 0: self.data[self.cursor]['edit_caption'] = edit_caption def set_caption(self, edit_caption=None): if self.cursor >= 0: self.data[self.cursor]['caption'] = self.data[self.cursor][ 'edit_caption'] if edit_caption is None else edit_caption class Text2ImageDataCard(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'] = 'txt2img' self.components_name = Text2ImageDataCardName(language) if self.new_dataset: # 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() else: for da_idx, cur_data in enumerate(self.data): if 'edit_caption' not in cur_data: self.data[da_idx]['edit_caption'] = cur_data['caption'] if 'edit_image_path' not in cur_data: self.data[da_idx]['edit_image_path'] = cur_data[ 'image_path'] if 'edit_relative_path' not in cur_data: self.data[da_idx]['edit_relative_path'] = cur_data[ 'relative_path'] if 'edit_width' not in cur_data: self.data[da_idx]['edit_width'] = cur_data['width'] if 'edit_height' not in cur_data: self.data[da_idx]['edit_height'] = cur_data['height'] 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('images') or one_dir.endswith('images/'): 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('/', '') == 'images': 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 if (one_s_dir.endswith('.jpg') or one_s_dir.endswith('.jpeg') or one_s_dir.endswith('.png') or one_s_dir.endswith('.webp')): file_name, surfix = os.path.splitext(one_s_dir) txt_file = file_name + '.txt' if os.path.exists(txt_file): raw_list[one_s_dir] = txt_file else: raw_list[one_s_dir] = None elif one_dir.endswith('train.csv'): train_list = one_dir else: if (one_dir.endswith('.jpg') or one_dir.endswith('.jpeg') or one_dir.endswith('.png') or one_dir.endswith('.webp')): file_name, surfix = os.path.splitext(one_dir) txt_file = file_name + '.txt' if os.path.exists(txt_file): raw_list[one_dir] = txt_file else: raw_list[one_dir] = None 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( "images folder or train.csv doesn't exists, or nothing exists in your zip" ) new_file_folder = f'{local_dataset_folder}/images' 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) elif len(raw_list) > 0: raw_list = [[k, v] for k, v in raw_list.items()] for img_id, cur_image in enumerate(raw_list): image_name, surfix = os.path.splitext(cur_image[0]) if cur_image[1] is not None and os.path.exists(cur_image[1]): prompt = open(cur_image[1], 'r').read() else: prompt = image_name.split('/')[-1] try: os.rename( os.path.abspath(cur_image[0]), f'{new_file_folder}/{get_md5(cur_image[0])}{surfix}') raw_list[img_id] = [ os.path.join('images', f'{get_md5(cur_image[0])}{surfix}'), prompt ] except Exception as e: print(e) if not os.path.exists(new_file_folder): raise gr.Error(f'{str(res)}') new_train_list = f'{local_dataset_folder}/train.csv' if train_list is None or not os.path.exists(train_list): with open(new_train_list, 'w') as f: writer = csv.writer(f) writer.writerow(['Target:FILE', 'Prompt']) for cur_image, cur_prompt in raw_list: writer.writerow([cur_image, cur_prompt]) else: 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}/images/*'") _ = res.readlines() res = os.popen(f"rm -rf '{hit_dir}'") _ = res.readlines() except Exception: pass file_list = self.load_train_file(new_train_list, data_folder) return file_list def load_train_file(self, file_path, data_folder): base_folder = os.path.dirname(file_path) file_list = [] with open(file_path, 'r') as f: reader = csv.reader(f) for row in reader: image_path, prompt = row[0], row[1] if image_path == 'Target:FILE': continue local_image_path = os.path.join(base_folder, image_path) w, h = get_image_meta(local_image_path) file_list.append({ 'image_path': os.path.join(data_folder, image_path), 'relative_path': image_path, 'width': w, 'height': h, 'caption': prompt, 'prefix': '', 'edit_caption': prompt, 'edit_image_path': os.path.join(data_folder, image_path), 'edit_relative_path': image_path, 'edit_width': w, 'edit_height': h, }) return file_list def load_from_list(self, save_file, dataset_folder, local_dataset_folder): file_list = [] images_folder = os.path.join(local_dataset_folder, 'images') os.makedirs(images_folder, exist_ok=True) with FS.get_from(save_file) as local_path: all_remote_list, all_local_list = [], [] all_save_list = [] with open(local_path, 'r') as f: for line in tqdm(f): line = line.strip() if line == '': continue try: image_path, width, height, caption = line.split( '#;#', 3) except Exception: try: image_path, width, height, caption = line.split( ',', 3) except Exception: raise gr.Error( self.components_name.illegal_data_err1) is_legal, new_path, prefix = find_prefix(image_path) try: int(width), int(height) except Exception: raise gr.Error( self.components_name.illegal_data_err4.format( width, height)) if not is_legal: raise gr.Error( self.components_name.illegal_data_err5.format( image_path)) relative_path = os.path.join('images', image_path.split('/')[-1]) all_remote_list.append(new_path) all_local_list.append( os.path.join(local_dataset_folder, relative_path)) all_save_list.append( os.path.join(dataset_folder, relative_path)) file_list.append({ 'image_path': os.path.join(dataset_folder, relative_path), 'relative_path': relative_path, 'width': int(width), 'height': int(height), 'caption': caption, 'prefix': prefix, 'edit_caption': caption, 'edit_image_path': os.path.join(dataset_folder, image_path), 'edit_relative_path': image_path, 'edit_width': int(width), 'edit_height': int(height), }) cache_file_list = [] for idx, local_path in enumerate( FS.get_batch_objects_from(all_remote_list)): if local_path is None: raise gr.Error( self.components_name.illegal_data_err6.format( all_remote_list[idx])) _ = FS.put_object_from_local_file(local_path, all_local_list[idx]) cache_file_list.append(local_path) for local_path, target_path, flg in FS.put_batch_objects_to( cache_file_list, all_save_list): if not flg: raise gr.Error( self.components_name.illegal_data_err7.format(local_path)) if os.path.exists(local_path): try: os.remove(local_path) except Exception: pass 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']]) 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: is_flag, file_path = del_prefix(one_file['image_path'], prefix=one_file['prefix']) f.write('{}#;#{}#;#{}#;#{}\n'.format(file_path, one_file['width'], one_file['height'], one_file['caption'])) FS.put_object_from_local_file(self.local_save_file_list, self.save_file_list) def add_record(self, image, caption, **kwargs): local_work_dir = self.meta['local_work_dir'] work_dir = self.meta['work_dir'] save_folder = os.path.join(local_work_dir, 'images') os.makedirs(save_folder, exist_ok=True) w, h = image.size relative_path = os.path.join('images', f'{imagehash.phash(image)}.png') image_path = os.path.join(work_dir, relative_path) local_image_path = os.path.join(local_work_dir, relative_path) with FS.put_to(image_path) as local_path: image.save(local_path) image.save(local_image_path) FS.put_object_from_local_file(local_image_path, image_path) self.data.append({ 'image_path': image_path, 'relative_path': relative_path, 'width': w, 'height': h, 'caption': caption, 'prefix': '', 'edit_caption': caption, 'edit_image_path': image_path, 'edit_relative_path': relative_path, 'edit_width': w, 'edit_height': h }) self.set_cursor(len(self.meta['file_list']) - 1) self.update_dataset() return True def delete_record(self): if len(self) < 1: raise gr.Error(self.components_name.delete_err1) current_file = self.data.pop(self.cursor) 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 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 images '{self.dataset_name}/images' " 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()) if not FS.exists(zip_path): raise gr.Error(self.components_name.export_zip_err1) return local_zip class Image2ImageDataCard(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'] = 'txt2img' self.components_name = Text2ImageDataCardName(language) if self.new_dataset: # 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 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('images') or one_dir.endswith('images/'): 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('/', '') == 'images': 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 if (one_s_dir.endswith('.jpg') or one_s_dir.endswith('.jpeg') or one_s_dir.endswith('.png') or one_s_dir.endswith('.webp')): file_name, surfix = os.path.splitext(one_s_dir) txt_file = file_name + '.txt' if os.path.exists(txt_file): raw_list[one_s_dir] = txt_file else: raw_list[one_s_dir] = None elif one_dir.endswith('train.csv'): train_list = one_dir else: if (one_dir.endswith('.jpg') or one_dir.endswith('.jpeg') or one_dir.endswith('.png') or one_dir.endswith('.webp')): file_name, surfix = os.path.splitext(one_dir) txt_file = file_name + '.txt' if os.path.exists(txt_file): raw_list[one_dir] = txt_file else: raw_list[one_dir] = None 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( "images 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}/images' os.makedirs(new_file_folder, exist_ok=True) images_dict = {} if file_folder is not None: _ = FS.get_dir_to_local_dir(file_folder, new_file_folder) elif len(raw_list) > 0: raw_list = [[k, v] for k, v in raw_list.items()] for img_id, cur_image in enumerate(raw_list): image_name, surfix = os.path.splitext(cur_image[0]) try: os.rename( os.path.abspath(cur_image[0]), f'{new_file_folder}/{get_md5(cur_image[0])}{surfix}') images_dict[os.path.basename(image_name)] = os.path.join( 'images', f'{get_md5(cur_image[0])}{surfix}'), except Exception as e: print(e) 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}/images/*'") _ = res.readlines() res = os.popen(f"rm -rf '{hit_dir}'") _ = res.readlines() except Exception: pass file_list = self.load_train_file(new_train_list, data_folder, images_dict) return file_list def load_train_file(self, file_path, data_folder, images_dict={}): base_folder = os.path.dirname(file_path) file_list = [] with open(file_path, 'r') as f: reader = csv.reader(f) for row in reader: src_image_path, image_path, prompt = row[0], row[1], row[2] if image_path == 'Target:FILE': continue src_image_name, surfix = os.path.splitext(src_image_path) src_image_path = images_dict.get( src_image_path, os.path.basename(src_image_name)) local_src_image_path = os.path.join(base_folder, src_image_path) src_w, src_h = get_image_meta(local_src_image_path) image_name, surfix = os.path.splitext(image_path) image_path = images_dict.get(image_path, os.path.basename(image_name)) local_image_path = os.path.join(base_folder, image_path) w, h = get_image_meta(local_image_path) file_list.append({ 'image_path': os.path.join(data_folder, image_path), 'relative_path': image_path, 'width': w, 'height': h, 'src_image_path': os.path.join(data_folder, src_image_path), 'src_relative_path': src_image_path, 'src_width': src_w, 'src_height': src_h, 'caption': prompt, 'prefix': '', 'edit_caption': prompt }) return file_list def load_from_list(self, save_file, dataset_folder, local_dataset_folder): file_list = [] images_folder = os.path.join(local_dataset_folder, 'images') os.makedirs(images_folder, exist_ok=True) with FS.get_from(save_file) as local_path: all_remote_list, all_local_list = [], [] all_src_remote_list, all_src_local_list = [], [] all_save_list, all_src_save_list = [], [] with open(local_path, 'r') as f: for line in tqdm(f): line = line.strip() if line == '': continue try: image_path, width, height, src_image_path, src_width, src_height, caption = line.split( '#;#', 6) except Exception: try: image_path, width, height, src_image_path, src_width, src_height, caption = line.split( ',', 6) except Exception: raise gr.Error( self.components_name.illegal_data_err1) is_legal, new_path, prefix = find_prefix(image_path) try: int(width), int(height) except Exception: raise gr.Error( self.components_name.illegal_data_err4.format( width, height)) is_legal, new_src_path, src_prefix = find_prefix( src_image_path) try: int(src_width), int(src_height) except Exception: raise gr.Error( self.components_name.illegal_data_err4.format( src_width, src_height)) if not is_legal: raise gr.Error( self.components_name.illegal_data_err5.format( image_path + ' ' + src_image_path)) relative_path = os.path.join('images', image_path.split('/')[-1]) src_relative_path = os.path.join( 'images', src_image_path.split('/')[-1]) all_remote_list.append(new_path) all_local_list.append( os.path.join(local_dataset_folder, relative_path)) all_save_list.append( os.path.join(dataset_folder, relative_path)) all_src_remote_list.append(new_src_path) all_src_local_list.append( os.path.join(local_dataset_folder, src_relative_path)) all_src_save_list.append( os.path.join(dataset_folder, src_relative_path)) file_list.append({ 'image_path': os.path.join(dataset_folder, relative_path), 'src_image_path': os.path.join(dataset_folder, src_relative_path), 'relative_path': relative_path, 'src_relative_path': src_relative_path, 'width': int(width), 'height': int(height), 'src_width': int(src_width), 'src_height': int(src_height), 'caption': caption, 'prefix': prefix, 'edit_caption': caption }) cache_file_list = [] for idx, local_path in enumerate( FS.get_batch_objects_from(all_remote_list)): if local_path is None: raise gr.Error( self.components_name.illegal_data_err6.format( all_remote_list[idx])) _ = FS.put_object_from_local_file(local_path, all_local_list[idx]) cache_file_list.append(local_path) for local_path, target_path, flg in FS.put_batch_objects_to( cache_file_list, all_save_list): if not flg: raise gr.Error( self.components_name.illegal_data_err7.format(local_path)) if os.path.exists(local_path): try: os.remove(local_path) except Exception: pass cache_src_file_list = [] for idx, local_path in enumerate( FS.get_batch_objects_from(all_src_remote_list)): if local_path is None: raise gr.Error( self.components_name.illegal_data_err6.format( all_src_remote_list[idx])) _ = FS.put_object_from_local_file(local_path, all_src_local_list[idx]) cache_src_file_list.append(local_path) for local_path, target_path, flg in FS.put_batch_objects_to( cache_src_file_list, all_src_save_list): if not flg: raise gr.Error( self.components_name.illegal_data_err7.format(local_path)) if os.path.exists(local_path): try: os.remove(local_path) except Exception: pass 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', 'Source:FILE', 'Prompt']) for one_file in file_list: relative_file = one_file['relative_path'] if relative_file.startswith('/'): relative_file = relative_file[1:] src_relative_file = one_file['src_relative_path'] if src_relative_file.startswith('/'): src_relative_file = src_relative_file[1:] writer.writerow( [relative_file, src_relative_file, one_file['caption']]) 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: is_flag, file_path = del_prefix(one_file['image_path'], prefix=one_file['prefix']) is_flag, src_file_path = del_prefix(one_file['src_image_path'], prefix=one_file['prefix']) f.write('{}#;#{}#;#{}#;#{}#;#{}#;#{}#;#{}\n'.format( file_path, one_file['width'], one_file['height'], src_file_path, one_file['src_width'], one_file['src_height'], one_file['caption'])) FS.put_object_from_local_file(self.local_save_file_list, self.save_file_list) def add_record(self, images, caption, **kwargs): local_work_dir = self.meta['local_work_dir'] work_dir = self.meta['work_dir'] save_folder = os.path.join(local_work_dir, 'images') os.makedirs(save_folder, exist_ok=True) image, src_image = images w, h = image.size relative_path = os.path.join('images', f'{imagehash.phash(image)}.png') image_path = os.path.join(work_dir, relative_path) local_image_path = os.path.join(local_work_dir, relative_path) image.save(local_image_path) FS.put_object_from_local_file(local_image_path, image_path) src_w, src_h = src_image.size src_relative_path = os.path.join('images', f'{imagehash.phash(src_image)}.png') src_image_path = os.path.join(work_dir, src_relative_path) local_src_image_path = os.path.join(local_work_dir, src_relative_path) image.save(local_src_image_path) FS.put_object_from_local_file(local_src_image_path, src_image_path) self.data.append({ 'image_path': image_path, 'relative_path': relative_path, 'width': w, 'height': h, 'src_image_path': src_image_path, 'src_relative_path': src_relative_path, 'src_width': src_w, 'src_height': src_h, 'caption': caption, 'prefix': '', 'edit_caption': caption }) self.set_cursor(len(self.meta['file_list']) - 1) self.update_dataset() return True def delete_record(self): if len(self) < 1: raise gr.Error(self.components_name.delete_err1) current_file = self.data.pop(self.cursor) 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 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 images '{self.dataset_name}/images' " 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()) if not FS.exists(zip_path): raise gr.Error(self.components_name.export_zip_err1) return local_zip