mingyuan/MotionDiffuse

1import os2import numpy as np3# import cv24from PIL import Image5from utils import paramUtil6import math7import time8import matplotlib.pyplot as plt9from scipy.ndimage import gaussian_filter10 11 12def mkdir(path):13    if not os.path.exists(path):14        os.makedirs(path)15 16COLORS = [[255, 0, 0], [255, 85, 0], [255, 170, 0], [255, 255, 0], [170, 255, 0], [85, 255, 0], [0, 255, 0],17          [0, 255, 85], [0, 255, 170], [0, 255, 255], [0, 170, 255], [0, 85, 255], [0, 0, 255], [85, 0, 255],18          [170, 0, 255], [255, 0, 255], [255, 0, 170], [255, 0, 85]]19 20MISSING_VALUE = -121 22def save_image(image_numpy, image_path):23    img_pil = Image.fromarray(image_numpy)24    img_pil.save(image_path)25 26 27def save_logfile(log_loss, save_path):28    with open(save_path, 'wt') as f:29        for k, v in log_loss.items():30            w_line = k31            for digit in v:32                w_line += ' %.3f' % digit33            f.write(w_line + '\n')34 35 36def print_current_loss(start_time, niter_state, losses, epoch=None, inner_iter=None):37 38    def as_minutes(s):39        m = math.floor(s / 60)40        s -= m * 6041        return '%dm %ds' % (m, s)42 43    def time_since(since, percent):44        now = time.time()45        s = now - since46        es = s / percent47        rs = es - s48        return '%s (- %s)' % (as_minutes(s), as_minutes(rs))49 50    if epoch is not None:51        print('epoch: %3d niter: %6d  inner_iter: %4d' % (epoch, niter_state, inner_iter), end=" ")52 53    now = time.time()54    message = '%s'%(as_minutes(now - start_time))55 56    for k, v in losses.items():57        message += ' %s: %.4f ' % (k, v)58    print(message)59 60 61def compose_gif_img_list(img_list, fp_out, duration):62    img, *imgs = [Image.fromarray(np.array(image)) for image in img_list]63    img.save(fp=fp_out, format='GIF', append_images=imgs, optimize=False,64             save_all=True, loop=0, duration=duration)65 66 67def save_images(visuals, image_path):68    if not os.path.exists(image_path):69        os.makedirs(image_path)70 71    for i, (label, img_numpy) in enumerate(visuals.items()):72        img_name = '%d_%s.jpg' % (i, label)73        save_path = os.path.join(image_path, img_name)74        save_image(img_numpy, save_path)75 76 77def save_images_test(visuals, image_path, from_name, to_name):78    if not os.path.exists(image_path):79        os.makedirs(image_path)80 81    for i, (label, img_numpy) in enumerate(visuals.items()):82        img_name = "%s_%s_%s" % (from_name, to_name, label)83        save_path = os.path.join(image_path, img_name)84        save_image(img_numpy, save_path)85 86 87def compose_and_save_img(img_list, save_dir, img_name, col=4, row=1, img_size=(256, 200)):88    # print(col, row)89    compose_img = compose_image(img_list, col, row, img_size)90    if not os.path.exists(save_dir):91        os.makedirs(save_dir)92    img_path = os.path.join(save_dir, img_name)93    # print(img_path)94    compose_img.save(img_path)95 96 97def compose_image(img_list, col, row, img_size):98    to_image = Image.new('RGB', (col * img_size[0], row * img_size[1]))99    for y in range(0, row):100        for x in range(0, col):101            from_img = Image.fromarray(img_list[y * col + x])102            # print((x * img_size[0], y*img_size[1],103            #                           (x + 1) * img_size[0], (y + 1) * img_size[1]))104            paste_area = (x * img_size[0], y*img_size[1],105                                      (x + 1) * img_size[0], (y + 1) * img_size[1])106            to_image.paste(from_img, paste_area)107            # to_image[y*img_size[1]:(y + 1) * img_size[1], x * img_size[0] :(x + 1) * img_size[0]] = from_img108    return to_image109 110 111def list_cut_average(ll, intervals):112    if intervals == 1:113        return ll114 115    bins = math.ceil(len(ll) * 1.0 / intervals)116    ll_new = []117    for i in range(bins):118        l_low = intervals * i119        l_high = l_low + intervals120        l_high = l_high if l_high < len(ll) else len(ll)121        ll_new.append(np.mean(ll[l_low:l_high]))122    return ll_new123 124 125def motion_temporal_filter(motion, sigma=1):126    motion = motion.reshape(motion.shape[0], -1)127    # print(motion.shape)128    for i in range(motion.shape[1]):129        motion[:, i] = gaussian_filter(motion[:, i], sigma=sigma, mode="nearest")130    return motion.reshape(motion.shape[0], -1, 3)131 132