From 41c045f476c8d9996f58e4b30c22edfd6f8b4e05 Mon Sep 17 00:00:00 2001 From: Isi <86603298+Isi-dev@users.noreply.github.com> Date: Mon, 5 Aug 2024 20:14:22 +0100 Subject: [PATCH] Delete test_func directory --- test_func/save_targer_keys.py | 219 ------------------------ test_func/test_EndDec.py | 193 --------------------- test_func/test_dataset.py | 307 ---------------------------------- test_func/test_models.py | 115 ------------- test_func/test_save_video.py | 51 ------ 5 files changed, 885 deletions(-) delete mode 100644 test_func/save_targer_keys.py delete mode 100644 test_func/test_EndDec.py delete mode 100644 test_func/test_dataset.py delete mode 100644 test_func/test_models.py delete mode 100644 test_func/test_save_video.py diff --git a/test_func/save_targer_keys.py b/test_func/save_targer_keys.py deleted file mode 100644 index 5aa322f..0000000 --- a/test_func/save_targer_keys.py +++ /dev/null @@ -1,219 +0,0 @@ -<<<<<<< HEAD -import os -import sys -import json -import torch -import imageio -import numpy as np -import os.path as osp -sys.path.insert(0, '/'.join(osp.realpath(__file__).split('/')[:-2])) -from thop import profile -from ptflops import get_model_complexity_info - -import artist.data as data -from tools.modules.config import cfg -from tools.modules.unet.util import * -from utils.config import Config as pConfig -from utils.registry_class import ENGINE, MODEL - - -def save_temporal_key(): - cfg_update = pConfig(load=True) - - for k, v in cfg_update.cfg_dict.items(): - if isinstance(v, dict) and k in cfg: - cfg[k].update(v) - else: - cfg[k] = v - - model = MODEL.build(cfg.UNet) - - temp_name = '' - temp_key_list = [] - spth = 'workspace/module_list/UNetSD_I2V_vs_Text_temporal_key_list.json' - for name, module in model.named_modules(): - if isinstance(module, (TemporalTransformer, TemporalTransformer_attemask, TemporalAttentionBlock, TemporalAttentionMultiBlock, TemporalConvBlock_v2, TemporalConvBlock)): - temp_name = name - print(f'Model: {name}') - elif isinstance(module, (ResidualBlock, ResBlock, SpatialTransformer, Upsample, Downsample)): - temp_name = '' - - if hasattr(module, 'weight'): - if temp_name != '' and (temp_name in name): - temp_key_list.append(name) - print(f'{name}') - # print(name) - - save_module_list = [] - for k, p in model.named_parameters(): - for item in temp_key_list: - if item in k: - print(f'{item} --> {k}') - save_module_list.append(k) - - print(int(sum(p.numel() for k, p in model.named_parameters()) / (1024 ** 2)), 'M parameters') - - # spth = 'workspace/module_list/{}' - json.dump(save_module_list, open(spth, 'w')) - a = 0 - - -def save_spatial_key(): - cfg_update = pConfig(load=True) - - for k, v in cfg_update.cfg_dict.items(): - if isinstance(v, dict) and k in cfg: - cfg[k].update(v) - else: - cfg[k] = v - - model = MODEL.build(cfg.UNet) - temp_name = '' - temp_key_list = [] - spth = 'workspace/module_list/UNetSD_I2V_HQ_P_spatial_key_list.json' - for name, module in model.named_modules(): - if isinstance(module, (ResidualBlock, ResBlock, SpatialTransformer, Upsample, Downsample)): - temp_name = name - print(f'Model: {name}') - elif isinstance(module, (TemporalTransformer, TemporalTransformer_attemask, TemporalAttentionBlock, TemporalAttentionMultiBlock, TemporalConvBlock_v2, TemporalConvBlock)): - temp_name = '' - - if hasattr(module, 'weight'): - if temp_name != '' and (temp_name in name): - temp_key_list.append(name) - print(f'{name}') - # print(name) - - save_module_list = [] - for k, p in model.named_parameters(): - for item in temp_key_list: - if item in k: - print(f'{item} --> {k}') - save_module_list.append(k) - - print(int(sum(p.numel() for k, p in model.named_parameters()) / (1024 ** 2)), 'M parameters') - - # spth = 'workspace/module_list/{}' - json.dump(save_module_list, open(spth, 'w')) - a = 0 - - -if __name__ == '__main__': - # save_temporal_key() - save_spatial_key() - - - -# print([k for (k, _) in self.input_blocks.named_parameters()]) - - -======= -import os -import sys -import json -import torch -import imageio -import numpy as np -import os.path as osp -sys.path.insert(0, '/'.join(osp.realpath(__file__).split('/')[:-2])) -from thop import profile -from ptflops import get_model_complexity_info - -import artist.data as data -from tools.modules.config import cfg -from tools.modules.unet.util import * -from utils.config import Config as pConfig -from utils.registry_class import ENGINE, MODEL - - -def save_temporal_key(): - cfg_update = pConfig(load=True) - - for k, v in cfg_update.cfg_dict.items(): - if isinstance(v, dict) and k in cfg: - cfg[k].update(v) - else: - cfg[k] = v - - model = MODEL.build(cfg.UNet) - - temp_name = '' - temp_key_list = [] - spth = 'workspace/module_list/UNetSD_I2V_vs_Text_temporal_key_list.json' - for name, module in model.named_modules(): - if isinstance(module, (TemporalTransformer, TemporalTransformer_attemask, TemporalAttentionBlock, TemporalAttentionMultiBlock, TemporalConvBlock_v2, TemporalConvBlock)): - temp_name = name - print(f'Model: {name}') - elif isinstance(module, (ResidualBlock, ResBlock, SpatialTransformer, Upsample, Downsample)): - temp_name = '' - - if hasattr(module, 'weight'): - if temp_name != '' and (temp_name in name): - temp_key_list.append(name) - print(f'{name}') - # print(name) - - save_module_list = [] - for k, p in model.named_parameters(): - for item in temp_key_list: - if item in k: - print(f'{item} --> {k}') - save_module_list.append(k) - - print(int(sum(p.numel() for k, p in model.named_parameters()) / (1024 ** 2)), 'M parameters') - - # spth = 'workspace/module_list/{}' - json.dump(save_module_list, open(spth, 'w')) - a = 0 - - -def save_spatial_key(): - cfg_update = pConfig(load=True) - - for k, v in cfg_update.cfg_dict.items(): - if isinstance(v, dict) and k in cfg: - cfg[k].update(v) - else: - cfg[k] = v - - model = MODEL.build(cfg.UNet) - temp_name = '' - temp_key_list = [] - spth = 'workspace/module_list/UNetSD_I2V_HQ_P_spatial_key_list.json' - for name, module in model.named_modules(): - if isinstance(module, (ResidualBlock, ResBlock, SpatialTransformer, Upsample, Downsample)): - temp_name = name - print(f'Model: {name}') - elif isinstance(module, (TemporalTransformer, TemporalTransformer_attemask, TemporalAttentionBlock, TemporalAttentionMultiBlock, TemporalConvBlock_v2, TemporalConvBlock)): - temp_name = '' - - if hasattr(module, 'weight'): - if temp_name != '' and (temp_name in name): - temp_key_list.append(name) - print(f'{name}') - # print(name) - - save_module_list = [] - for k, p in model.named_parameters(): - for item in temp_key_list: - if item in k: - print(f'{item} --> {k}') - save_module_list.append(k) - - print(int(sum(p.numel() for k, p in model.named_parameters()) / (1024 ** 2)), 'M parameters') - - # spth = 'workspace/module_list/{}' - json.dump(save_module_list, open(spth, 'w')) - a = 0 - - -if __name__ == '__main__': - # save_temporal_key() - save_spatial_key() - - - -# print([k for (k, _) in self.input_blocks.named_parameters()]) - - ->>>>>>> 626e7afc02230297b6f553675ea1c32c29971314 diff --git a/test_func/test_EndDec.py b/test_func/test_EndDec.py deleted file mode 100644 index 5b18d97..0000000 --- a/test_func/test_EndDec.py +++ /dev/null @@ -1,193 +0,0 @@ -<<<<<<< HEAD -import os -import sys -import torch -import imageio -import numpy as np -import os.path as osp -sys.path.insert(0, '/'.join(osp.realpath(__file__).split('/')[:-2])) -from PIL import Image, ImageDraw, ImageFont - -from einops import rearrange - -from tools import * -import utils.transforms as data -from utils.seed import setup_seed -from tools.modules.config import cfg -from utils.config import Config as pConfig -from utils.registry_class import ENGINE, DATASETS, AUTO_ENCODER - - -def test_enc_dec(gpu=0): - setup_seed(0) - cfg_update = pConfig(load=True) - - for k, v in cfg_update.cfg_dict.items(): - if isinstance(v, dict) and k in cfg: - cfg[k].update(v) - else: - cfg[k] = v - - save_dir = os.path.join('workspace/test_data/autoencoder', cfg.auto_encoder['type']) - os.system('rm -rf %s' % (save_dir)) - os.makedirs(save_dir, exist_ok=True) - - train_trans = data.Compose([ - data.CenterCropWide(size=cfg.resolution), - data.ToTensor(), - data.Normalize(mean=cfg.mean, std=cfg.std)]) - - vit_trans = data.Compose([ - data.CenterCropWide(size=(cfg.resolution[0], cfg.resolution[0])) if cfg.resolution[0]>cfg.vit_resolution[0] else data.CenterCropWide(size=cfg.vit_resolution), - data.Resize(cfg.vit_resolution), - data.ToTensor(), - data.Normalize(mean=cfg.vit_mean, std=cfg.vit_std)]) - - video_mean = torch.tensor(cfg.mean).view(1, -1, 1, 1) #n c f h w - video_std = torch.tensor(cfg.std).view(1, -1, 1, 1) #n c f h w - - txt_size = cfg.resolution[1] - nc = int(38 * (txt_size / 256)) - font = ImageFont.truetype('data/font/DejaVuSans.ttf', size=13) - - dataset = DATASETS.build(cfg.vid_dataset, sample_fps=4, transforms=train_trans, vit_transforms=vit_trans) - print('There are %d videos' % (len(dataset))) - - autoencoder = AUTO_ENCODER.build(cfg.auto_encoder) - autoencoder.eval() # freeze - for param in autoencoder.parameters(): - param.requires_grad = False - autoencoder.to(gpu) - for idx, item in enumerate(dataset): - local_path = os.path.join(save_dir, '%04d.mp4' % idx) - # ref_frame, video_data, caption = item - ref_frame, vit_frame, video_data = item[:3] - video_data = video_data.to(gpu) - - image_list = [] - video_data_list = torch.chunk(video_data, video_data.shape[0]//cfg.chunk_size,dim=0) - with torch.no_grad(): - decode_data = [] - for chunk_data in video_data_list: - latent_z = autoencoder.encode_firsr_stage(chunk_data).detach() - # latent_z = get_first_stage_encoding(encoder_posterior).detach() - kwargs = {"timesteps": chunk_data.shape[0]} - recons_data = autoencoder.decode(latent_z, **kwargs) - - vis_data = torch.cat([chunk_data, recons_data], dim=2).cpu() - vis_data = vis_data.mul_(video_std).add_(video_mean) # 8x3x16x256x384 - vis_data = vis_data.cpu() - vis_data.clamp_(0, 1) - vis_data = vis_data.permute(0, 2, 3, 1) - vis_data = [(image.numpy() * 255).astype('uint8') for image in vis_data] - image_list.extend(vis_data) - - num_image = len(image_list) - frame_dir = os.path.join(save_dir, 'temp') - os.makedirs(frame_dir, exist_ok=True) - for idx in range(num_image): - tpth = os.path.join(frame_dir, '%04d.png' % (idx+1)) - cv2.imwrite(tpth, image_list[idx][:,:,::-1], [int(cv2.IMWRITE_JPEG_QUALITY), 100]) - cmd = f'ffmpeg -y -f image2 -loglevel quiet -framerate 8 -i {frame_dir}/%04d.png -vcodec libx264 -crf 17 -pix_fmt yuv420p {local_path}' - os.system(cmd); os.system(f'rm -rf {frame_dir}') - - -if __name__ == '__main__': - test_enc_dec() -======= -import os -import sys -import torch -import imageio -import numpy as np -import os.path as osp -sys.path.insert(0, '/'.join(osp.realpath(__file__).split('/')[:-2])) -from PIL import Image, ImageDraw, ImageFont - -from einops import rearrange - -from tools import * -import utils.transforms as data -from utils.seed import setup_seed -from tools.modules.config import cfg -from utils.config import Config as pConfig -from utils.registry_class import ENGINE, DATASETS, AUTO_ENCODER - - -def test_enc_dec(gpu=0): - setup_seed(0) - cfg_update = pConfig(load=True) - - for k, v in cfg_update.cfg_dict.items(): - if isinstance(v, dict) and k in cfg: - cfg[k].update(v) - else: - cfg[k] = v - - save_dir = os.path.join('workspace/test_data/autoencoder', cfg.auto_encoder['type']) - os.system('rm -rf %s' % (save_dir)) - os.makedirs(save_dir, exist_ok=True) - - train_trans = data.Compose([ - data.CenterCropWide(size=cfg.resolution), - data.ToTensor(), - data.Normalize(mean=cfg.mean, std=cfg.std)]) - - vit_trans = data.Compose([ - data.CenterCropWide(size=(cfg.resolution[0], cfg.resolution[0])) if cfg.resolution[0]>cfg.vit_resolution[0] else data.CenterCropWide(size=cfg.vit_resolution), - data.Resize(cfg.vit_resolution), - data.ToTensor(), - data.Normalize(mean=cfg.vit_mean, std=cfg.vit_std)]) - - video_mean = torch.tensor(cfg.mean).view(1, -1, 1, 1) #n c f h w - video_std = torch.tensor(cfg.std).view(1, -1, 1, 1) #n c f h w - - txt_size = cfg.resolution[1] - nc = int(38 * (txt_size / 256)) - font = ImageFont.truetype('data/font/DejaVuSans.ttf', size=13) - - dataset = DATASETS.build(cfg.vid_dataset, sample_fps=4, transforms=train_trans, vit_transforms=vit_trans) - print('There are %d videos' % (len(dataset))) - - autoencoder = AUTO_ENCODER.build(cfg.auto_encoder) - autoencoder.eval() # freeze - for param in autoencoder.parameters(): - param.requires_grad = False - autoencoder.to(gpu) - for idx, item in enumerate(dataset): - local_path = os.path.join(save_dir, '%04d.mp4' % idx) - # ref_frame, video_data, caption = item - ref_frame, vit_frame, video_data = item[:3] - video_data = video_data.to(gpu) - - image_list = [] - video_data_list = torch.chunk(video_data, video_data.shape[0]//cfg.chunk_size,dim=0) - with torch.no_grad(): - decode_data = [] - for chunk_data in video_data_list: - latent_z = autoencoder.encode_firsr_stage(chunk_data).detach() - # latent_z = get_first_stage_encoding(encoder_posterior).detach() - kwargs = {"timesteps": chunk_data.shape[0]} - recons_data = autoencoder.decode(latent_z, **kwargs) - - vis_data = torch.cat([chunk_data, recons_data], dim=2).cpu() - vis_data = vis_data.mul_(video_std).add_(video_mean) # 8x3x16x256x384 - vis_data = vis_data.cpu() - vis_data.clamp_(0, 1) - vis_data = vis_data.permute(0, 2, 3, 1) - vis_data = [(image.numpy() * 255).astype('uint8') for image in vis_data] - image_list.extend(vis_data) - - num_image = len(image_list) - frame_dir = os.path.join(save_dir, 'temp') - os.makedirs(frame_dir, exist_ok=True) - for idx in range(num_image): - tpth = os.path.join(frame_dir, '%04d.png' % (idx+1)) - cv2.imwrite(tpth, image_list[idx][:,:,::-1], [int(cv2.IMWRITE_JPEG_QUALITY), 100]) - cmd = f'ffmpeg -y -f image2 -loglevel quiet -framerate 8 -i {frame_dir}/%04d.png -vcodec libx264 -crf 17 -pix_fmt yuv420p {local_path}' - os.system(cmd); os.system(f'rm -rf {frame_dir}') - - -if __name__ == '__main__': - test_enc_dec() ->>>>>>> 626e7afc02230297b6f553675ea1c32c29971314 diff --git a/test_func/test_dataset.py b/test_func/test_dataset.py deleted file mode 100644 index 1f0745d..0000000 --- a/test_func/test_dataset.py +++ /dev/null @@ -1,307 +0,0 @@ -<<<<<<< HEAD -import os -import sys -import imageio -import numpy as np -import os.path as osp -sys.path.insert(0, '/'.join(osp.realpath(__file__).split('/')[:-2])) -from PIL import Image, ImageDraw, ImageFont -import torchvision.transforms as T - -import utils.transforms as data -from tools.modules.config import cfg -from utils.config import Config as pConfig -from utils.registry_class import ENGINE, DATASETS - -from tools import * - -def test_video_dataset(): - cfg_update = pConfig(load=True) - - for k, v in cfg_update.cfg_dict.items(): - if isinstance(v, dict) and k in cfg: - cfg[k].update(v) - else: - cfg[k] = v - - exp_name = os.path.basename(cfg.cfg_file).split('.')[0] - save_dir = os.path.join('workspace', 'test_data/datasets', cfg.vid_dataset['type'], exp_name) - os.system('rm -rf %s' % (save_dir)) - os.makedirs(save_dir, exist_ok=True) - - train_trans = data.Compose([ - data.CenterCropWide(size=cfg.resolution), - data.ToTensor(), - data.Normalize(mean=cfg.mean, std=cfg.std)]) - vit_trans = T.Compose([ - data.CenterCropWide(cfg.vit_resolution), - T.ToTensor(), - T.Normalize(mean=cfg.vit_mean, std=cfg.vit_std)]) - - video_mean = torch.tensor(cfg.mean).view(1, -1, 1, 1) #n c f h w - video_std = torch.tensor(cfg.std).view(1, -1, 1, 1) #n c f h w - - img_mean = torch.tensor(cfg.mean).view(-1, 1, 1) # c f h w - img_std = torch.tensor(cfg.std).view(-1, 1, 1) # c f h w - - vit_mean = torch.tensor(cfg.vit_mean).view(-1, 1, 1) # c f h w - vit_std = torch.tensor(cfg.vit_std).view(-1, 1, 1) # c f h w - - txt_size = cfg.resolution[1] - nc = int(38 * (txt_size / 256)) - font = ImageFont.truetype('data/font/DejaVuSans.ttf', size=13) - - dataset = DATASETS.build(cfg.vid_dataset, sample_fps=cfg.sample_fps[0], transforms=train_trans, vit_transforms=vit_trans) - print('There are %d videos' % (len(dataset))) - for idx, item in enumerate(dataset): - ref_frame, vit_frame, video_data, caption, video_key = item - - video_data = video_data.mul_(video_std).add_(video_mean) - video_data.clamp_(0, 1) - video_data = video_data.permute(0, 2, 3, 1) - video_data = [(image.numpy() * 255).astype('uint8') for image in video_data] - - # Single Image - ref_frame = ref_frame.mul_(img_mean).add_(img_std) - ref_frame.clamp_(0, 1) - ref_frame = ref_frame.permute(1, 2, 0) - ref_frame = (ref_frame.numpy() * 255).astype('uint8') - - # Text image - txt_img = Image.new("RGB", (txt_size, txt_size), color="white") - draw = ImageDraw.Draw(txt_img) - lines = "\n".join(caption[start:start + nc] for start in range(0, len(caption), nc)) - draw.text((0, 0), lines, fill="black", font=font) - txt_img = np.array(txt_img) - - video_data = [np.concatenate([ref_frame, u, txt_img], axis=1) for u in video_data] - spath = os.path.join(save_dir, '%04d.gif' % (idx)) - imageio.mimwrite(spath, video_data, fps =8) - - # if idx > 100: break - - -def test_vit_image(test_video_flag=True): - cfg_update = pConfig(load=True) - - for k, v in cfg_update.cfg_dict.items(): - if isinstance(v, dict) and k in cfg: - cfg[k].update(v) - else: - cfg[k] = v - - exp_name = os.path.basename(cfg.cfg_file).split('.')[0] - save_dir = os.path.join('workspace', 'test_data/datasets', cfg.img_dataset['type'], exp_name) - os.system('rm -rf %s' % (save_dir)) - os.makedirs(save_dir, exist_ok=True) - - train_trans = data.Compose([ - data.CenterCropWide(size=cfg.resolution), - data.ToTensor(), - data.Normalize(mean=cfg.mean, std=cfg.std)]) - vit_trans = data.Compose([ - data.CenterCropWide(cfg.resolution), - data.Resize(cfg.vit_resolution), - data.ToTensor(), - data.Normalize(mean=cfg.vit_mean, std=cfg.vit_std)]) - - img_mean = torch.tensor(cfg.mean).view(-1, 1, 1) # c f h w - img_std = torch.tensor(cfg.std).view(-1, 1, 1) # c f h w - - vit_mean = torch.tensor(cfg.vit_mean).view(-1, 1, 1) # c f h w - vit_std = torch.tensor(cfg.vit_std).view(-1, 1, 1) # c f h w - - txt_size = cfg.resolution[1] - nc = int(38 * (txt_size / 256)) - font = ImageFont.truetype('artist/font/DejaVuSans.ttf', size=13) - - dataset = DATASETS.build(cfg.img_dataset, transforms=train_trans, vit_transforms=vit_trans) - print('There are %d videos' % (len(dataset))) - for idx, item in enumerate(dataset): - ref_frame, vit_frame, video_data, caption, video_key = item - video_data = video_data.mul_(img_std).add_(img_mean) - video_data.clamp_(0, 1) - video_data = video_data.permute(0, 2, 3, 1) - video_data = [(image.numpy() * 255).astype('uint8') for image in video_data] - - # Single Image - vit_frame = vit_frame.mul_(vit_std).add_(vit_mean) - vit_frame.clamp_(0, 1) - vit_frame = vit_frame.permute(1, 2, 0) - vit_frame = (vit_frame.numpy() * 255).astype('uint8') - - zero_frame = np.zeros((cfg.resolution[1], cfg.resolution[1], 3), dtype=np.uint8) - zero_frame[:vit_frame.shape[0], :vit_frame.shape[1], :] = vit_frame - - # Text image - txt_img = Image.new("RGB", (txt_size, txt_size), color="white") - draw = ImageDraw.Draw(txt_img) - lines = "\n".join(caption[start:start + nc] for start in range(0, len(caption), nc)) - draw.text((0, 0), lines, fill="black", font=font) - txt_img = np.array(txt_img) - - video_data = [np.concatenate([zero_frame, u, txt_img], axis=1) for u in video_data] - spath = os.path.join(save_dir, '%04d.gif' % (idx)) - imageio.mimwrite(spath, video_data, fps =8) - - # if idx > 100: break - - -if __name__ == '__main__': - # test_video_dataset() - test_vit_image() - -======= -import os -import sys -import imageio -import numpy as np -import os.path as osp -sys.path.insert(0, '/'.join(osp.realpath(__file__).split('/')[:-2])) -from PIL import Image, ImageDraw, ImageFont -import torchvision.transforms as T - -import utils.transforms as data -from tools.modules.config import cfg -from utils.config import Config as pConfig -from utils.registry_class import ENGINE, DATASETS - -from tools import * - -def test_video_dataset(): - cfg_update = pConfig(load=True) - - for k, v in cfg_update.cfg_dict.items(): - if isinstance(v, dict) and k in cfg: - cfg[k].update(v) - else: - cfg[k] = v - - exp_name = os.path.basename(cfg.cfg_file).split('.')[0] - save_dir = os.path.join('workspace', 'test_data/datasets', cfg.vid_dataset['type'], exp_name) - os.system('rm -rf %s' % (save_dir)) - os.makedirs(save_dir, exist_ok=True) - - train_trans = data.Compose([ - data.CenterCropWide(size=cfg.resolution), - data.ToTensor(), - data.Normalize(mean=cfg.mean, std=cfg.std)]) - vit_trans = T.Compose([ - data.CenterCropWide(cfg.vit_resolution), - T.ToTensor(), - T.Normalize(mean=cfg.vit_mean, std=cfg.vit_std)]) - - video_mean = torch.tensor(cfg.mean).view(1, -1, 1, 1) #n c f h w - video_std = torch.tensor(cfg.std).view(1, -1, 1, 1) #n c f h w - - img_mean = torch.tensor(cfg.mean).view(-1, 1, 1) # c f h w - img_std = torch.tensor(cfg.std).view(-1, 1, 1) # c f h w - - vit_mean = torch.tensor(cfg.vit_mean).view(-1, 1, 1) # c f h w - vit_std = torch.tensor(cfg.vit_std).view(-1, 1, 1) # c f h w - - txt_size = cfg.resolution[1] - nc = int(38 * (txt_size / 256)) - font = ImageFont.truetype('data/font/DejaVuSans.ttf', size=13) - - dataset = DATASETS.build(cfg.vid_dataset, sample_fps=cfg.sample_fps[0], transforms=train_trans, vit_transforms=vit_trans) - print('There are %d videos' % (len(dataset))) - for idx, item in enumerate(dataset): - ref_frame, vit_frame, video_data, caption, video_key = item - - video_data = video_data.mul_(video_std).add_(video_mean) - video_data.clamp_(0, 1) - video_data = video_data.permute(0, 2, 3, 1) - video_data = [(image.numpy() * 255).astype('uint8') for image in video_data] - - # Single Image - ref_frame = ref_frame.mul_(img_mean).add_(img_std) - ref_frame.clamp_(0, 1) - ref_frame = ref_frame.permute(1, 2, 0) - ref_frame = (ref_frame.numpy() * 255).astype('uint8') - - # Text image - txt_img = Image.new("RGB", (txt_size, txt_size), color="white") - draw = ImageDraw.Draw(txt_img) - lines = "\n".join(caption[start:start + nc] for start in range(0, len(caption), nc)) - draw.text((0, 0), lines, fill="black", font=font) - txt_img = np.array(txt_img) - - video_data = [np.concatenate([ref_frame, u, txt_img], axis=1) for u in video_data] - spath = os.path.join(save_dir, '%04d.gif' % (idx)) - imageio.mimwrite(spath, video_data, fps =8) - - # if idx > 100: break - - -def test_vit_image(test_video_flag=True): - cfg_update = pConfig(load=True) - - for k, v in cfg_update.cfg_dict.items(): - if isinstance(v, dict) and k in cfg: - cfg[k].update(v) - else: - cfg[k] = v - - exp_name = os.path.basename(cfg.cfg_file).split('.')[0] - save_dir = os.path.join('workspace', 'test_data/datasets', cfg.img_dataset['type'], exp_name) - os.system('rm -rf %s' % (save_dir)) - os.makedirs(save_dir, exist_ok=True) - - train_trans = data.Compose([ - data.CenterCropWide(size=cfg.resolution), - data.ToTensor(), - data.Normalize(mean=cfg.mean, std=cfg.std)]) - vit_trans = data.Compose([ - data.CenterCropWide(cfg.resolution), - data.Resize(cfg.vit_resolution), - data.ToTensor(), - data.Normalize(mean=cfg.vit_mean, std=cfg.vit_std)]) - - img_mean = torch.tensor(cfg.mean).view(-1, 1, 1) # c f h w - img_std = torch.tensor(cfg.std).view(-1, 1, 1) # c f h w - - vit_mean = torch.tensor(cfg.vit_mean).view(-1, 1, 1) # c f h w - vit_std = torch.tensor(cfg.vit_std).view(-1, 1, 1) # c f h w - - txt_size = cfg.resolution[1] - nc = int(38 * (txt_size / 256)) - font = ImageFont.truetype('artist/font/DejaVuSans.ttf', size=13) - - dataset = DATASETS.build(cfg.img_dataset, transforms=train_trans, vit_transforms=vit_trans) - print('There are %d videos' % (len(dataset))) - for idx, item in enumerate(dataset): - ref_frame, vit_frame, video_data, caption, video_key = item - video_data = video_data.mul_(img_std).add_(img_mean) - video_data.clamp_(0, 1) - video_data = video_data.permute(0, 2, 3, 1) - video_data = [(image.numpy() * 255).astype('uint8') for image in video_data] - - # Single Image - vit_frame = vit_frame.mul_(vit_std).add_(vit_mean) - vit_frame.clamp_(0, 1) - vit_frame = vit_frame.permute(1, 2, 0) - vit_frame = (vit_frame.numpy() * 255).astype('uint8') - - zero_frame = np.zeros((cfg.resolution[1], cfg.resolution[1], 3), dtype=np.uint8) - zero_frame[:vit_frame.shape[0], :vit_frame.shape[1], :] = vit_frame - - # Text image - txt_img = Image.new("RGB", (txt_size, txt_size), color="white") - draw = ImageDraw.Draw(txt_img) - lines = "\n".join(caption[start:start + nc] for start in range(0, len(caption), nc)) - draw.text((0, 0), lines, fill="black", font=font) - txt_img = np.array(txt_img) - - video_data = [np.concatenate([zero_frame, u, txt_img], axis=1) for u in video_data] - spath = os.path.join(save_dir, '%04d.gif' % (idx)) - imageio.mimwrite(spath, video_data, fps =8) - - # if idx > 100: break - - -if __name__ == '__main__': - # test_video_dataset() - test_vit_image() - ->>>>>>> 626e7afc02230297b6f553675ea1c32c29971314 diff --git a/test_func/test_models.py b/test_func/test_models.py deleted file mode 100644 index e5f37a7..0000000 --- a/test_func/test_models.py +++ /dev/null @@ -1,115 +0,0 @@ -<<<<<<< HEAD -import os -import sys -import torch -import imageio -import numpy as np -import os.path as osp -sys.path.insert(0, '/'.join(osp.realpath(__file__).split('/')[:-2])) -from thop import profile -from ptflops import get_model_complexity_info - -import artist.data as data -from tools.modules.config import cfg -from utils.config import Config as pConfig -from utils.registry_class import ENGINE, MODEL - - -def test_model(): - cfg_update = pConfig(load=True) - - for k, v in cfg_update.cfg_dict.items(): - if isinstance(v, dict) and k in cfg: - cfg[k].update(v) - else: - cfg[k] = v - - model = MODEL.build(cfg.UNet) - print(int(sum(p.numel() for k, p in model.named_parameters()) / (1024 ** 2)), 'M parameters') - - # state_dict = torch.load('cache/pretrain_model/jiuniu_0600000.pth', map_location='cpu') - # model.load_state_dict(state_dict, strict=False) - model = model.cuda() - - x = torch.Tensor(1, 4, 16, 32, 56).cuda() - t = torch.Tensor(1).cuda() - sims = torch.Tensor(1, 32).cuda() - fps = torch.Tensor([8]).cuda() - y = torch.Tensor(1, 1, 1024).cuda() - image = torch.Tensor(1, 3, 256, 448).cuda() - - ret = model(x=x, t=t, y=y, ori_img=image, sims=sims, fps=fps) - print('Out shape if {}'.format(ret.shape)) - - # flops, params = profile(model=model, inputs=(x, t, y, image, sims, fps)) - # print('Model: {:.2f} GFLOPs and {:.2f}M parameters'.format(flops/1e9, params/1e6)) - - def prepare_input(resolution): - return dict(x=[x, t, y, image, sims, fps]) - - flops, params = get_model_complexity_info(model, (1, 4, 16, 32, 56), - input_constructor = prepare_input, - as_strings=True, print_per_layer_stat=True) - print(' - Flops: ' + flops) - print(' - Params: ' + params) - -if __name__ == '__main__': - test_model() -======= -import os -import sys -import torch -import imageio -import numpy as np -import os.path as osp -sys.path.insert(0, '/'.join(osp.realpath(__file__).split('/')[:-2])) -from thop import profile -from ptflops import get_model_complexity_info - -import artist.data as data -from tools.modules.config import cfg -from utils.config import Config as pConfig -from utils.registry_class import ENGINE, MODEL - - -def test_model(): - cfg_update = pConfig(load=True) - - for k, v in cfg_update.cfg_dict.items(): - if isinstance(v, dict) and k in cfg: - cfg[k].update(v) - else: - cfg[k] = v - - model = MODEL.build(cfg.UNet) - print(int(sum(p.numel() for k, p in model.named_parameters()) / (1024 ** 2)), 'M parameters') - - # state_dict = torch.load('cache/pretrain_model/jiuniu_0600000.pth', map_location='cpu') - # model.load_state_dict(state_dict, strict=False) - model = model.cuda() - - x = torch.Tensor(1, 4, 16, 32, 56).cuda() - t = torch.Tensor(1).cuda() - sims = torch.Tensor(1, 32).cuda() - fps = torch.Tensor([8]).cuda() - y = torch.Tensor(1, 1, 1024).cuda() - image = torch.Tensor(1, 3, 256, 448).cuda() - - ret = model(x=x, t=t, y=y, ori_img=image, sims=sims, fps=fps) - print('Out shape if {}'.format(ret.shape)) - - # flops, params = profile(model=model, inputs=(x, t, y, image, sims, fps)) - # print('Model: {:.2f} GFLOPs and {:.2f}M parameters'.format(flops/1e9, params/1e6)) - - def prepare_input(resolution): - return dict(x=[x, t, y, image, sims, fps]) - - flops, params = get_model_complexity_info(model, (1, 4, 16, 32, 56), - input_constructor = prepare_input, - as_strings=True, print_per_layer_stat=True) - print(' - Flops: ' + flops) - print(' - Params: ' + params) - -if __name__ == '__main__': - test_model() ->>>>>>> 626e7afc02230297b6f553675ea1c32c29971314 diff --git a/test_func/test_save_video.py b/test_func/test_save_video.py deleted file mode 100644 index af687ec..0000000 --- a/test_func/test_save_video.py +++ /dev/null @@ -1,51 +0,0 @@ -<<<<<<< HEAD -import numpy as np -import cv2 - -cap = cv2.VideoCapture('workspace/img_dir/tst.mp4') - -fourcc = cv2.VideoWriter_fourcc(*'H264') - -ret, frame = cap.read() -vid_size = frame.shape[:2][::-1] - -out = cv2.VideoWriter('workspace/img_dir/testwrite.mp4',fourcc, 8, vid_size) -out.write(frame) - -while(cap.isOpened()): - ret, frame = cap.read() - if not ret: break - out.write(frame) - - -cap.release() -out.release() - - - -======= -import numpy as np -import cv2 - -cap = cv2.VideoCapture('workspace/img_dir/tst.mp4') - -fourcc = cv2.VideoWriter_fourcc(*'H264') - -ret, frame = cap.read() -vid_size = frame.shape[:2][::-1] - -out = cv2.VideoWriter('workspace/img_dir/testwrite.mp4',fourcc, 8, vid_size) -out.write(frame) - -while(cap.isOpened()): - ret, frame = cap.read() - if not ret: break - out.write(frame) - - -cap.release() -out.release() - - - ->>>>>>> 626e7afc02230297b6f553675ea1c32c29971314