Delete test_func directory
This commit is contained in:
Isi
@@ -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
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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
|
||||
Reference in New Issue
Block a user