Spaces:
Running
Running
| import os | |
| import numpy as np | |
| import pydiffvg | |
| import torch | |
| from scipy.optimize import curve_fit | |
| def get_svg_file(path): | |
| files = os.listdir(path) | |
| files = [f for f in files if "best.svg" in f] | |
| return files[0] | |
| def get_seed(filename): | |
| filename = filename[:-9] | |
| keyword = 'seed' | |
| before_keyword, keyword, after_keyword = filename.partition(keyword) | |
| return after_keyword | |
| def get_clip_loss(path, layer): | |
| path_config = path / "config.npy" | |
| config = np.load(path_config, allow_pickle=True)[()] | |
| loss_clip = np.array(config[f"loss_eval"]) | |
| best_iter = np.argsort(loss_clip)[0] | |
| loss_clip_layer = np.array(config[f"clip_vit_l{layer}_original_eval"]) | |
| return loss_clip, best_iter, loss_clip_layer | |
| def ratios_to_str(ratios): | |
| ratios_str = "" | |
| for r_ in ratios: | |
| r_str = f"{r_:.3f}" | |
| ratios_str += f"{float(r_str)}," | |
| ratios_str = ratios_str[:-1] | |
| return ratios_str | |
| def func(x, a, c, d): | |
| return a * np.exp(c * x) | |
| def func_inv(y, a, c, d): | |
| return np.log(y / a) * (1 / c) | |
| def get_func(ratios_rel, start_x, start_ys): | |
| target_ys = ratios_rel[start_ys:] | |
| x = np.linspace(start_x, start_x + len(target_ys) - 1, len(target_ys)) | |
| # calculate exponent | |
| popt, pcov = curve_fit(func, x, target_ys, maxfev=3000) | |
| return popt | |
| def get_clip_loss2(path, layer, object_or_background): | |
| path_config = path / "config.npy" | |
| config = np.load(path_config, allow_pickle=True)[()] | |
| loss_clip = np.array(config[f"loss_eval"]) | |
| best_iter = np.argsort(loss_clip)[0] | |
| loss_clip_layer = np.array(config[f"clip_vit_l{layer}_original_eval"]) | |
| if object_or_background == "object": | |
| loss_clip_layer4 = np.array(config[f"clip_vit_l4_original_eval"]) | |
| loss_clip_layer = 1 * loss_clip_layer4 + loss_clip_layer | |
| return best_iter, loss_clip_layer | |
| def get_ratios_dict(path_to_initial_sketches, folder_name_l, layer, im_name, object_or_background, step_size_l, | |
| num_ratios=8): | |
| # get the sketch of the given layer, and get L_clip_i | |
| svg_filename = get_svg_file(path_to_initial_sketches / folder_name_l) | |
| seed = get_seed(svg_filename) | |
| path_li = path_to_initial_sketches / folder_name_l / f"{folder_name_l}_seed{seed}" | |
| best_iter, loss_clip_layer = get_clip_loss2(path_li, layer, object_or_background) | |
| best_lclip_layer = loss_clip_layer[best_iter] | |
| r_1_k = 1 / best_lclip_layer | |
| # get the next ratios by jumping by 2 | |
| r_j_k = r_1_k | |
| ratios_k = [r_1_k] | |
| for j in range(4): | |
| r_j_k = r_j_k / 2 | |
| ratios_k.append(r_j_k) | |
| start_ys, start_x, end_x_addition = 0, 0, 0 | |
| popt = get_func(ratios_k, start_x=0, start_ys=0) # fit the function to ratios_k | |
| x_1_k = func_inv([r_1_k], *popt) | |
| step_size = step_size_l | |
| num_steps = num_ratios - start_x + end_x_addition | |
| start_ = x_1_k[0] | |
| end = num_steps * step_size | |
| # sample the function from the initial scaled r_1 with the corresponding step size | |
| new_xs_layer_l = np.linspace(start_, end - step_size + start_, num_steps) | |
| # print("new_xs_layer_l", new_xs_layer_l) | |
| ratios_li = func(new_xs_layer_l, *popt) | |
| ratios_str = ratios_to_str(ratios_li) | |
| xs_layer_l_str = ratios_to_str(new_xs_layer_l) | |
| print(f"layer {layer} r_1_k {r_1_k} \n new {ratios_str} \n x {xs_layer_l_str}\n") | |
| return ratios_str | |
| def read_svg(path_svg, multiply=0, resize_obj=False, params=None, opacity=1, device=None): | |
| pydiffvg.set_device(device) | |
| canvas_width, canvas_height, shapes, shape_groups = pydiffvg.svg_to_scene(path_svg) | |
| for group in shape_groups: | |
| group.stroke_color = torch.tensor([0, 0, 0, opacity]) | |
| if resize_obj and params: | |
| w, h = params["scale_w"], params["scale_h"] | |
| for path in shapes: | |
| path.points = path.points / canvas_width | |
| path.points = 2 * path.points - 1 | |
| path.points[:, 0] /= (w) # / canvas_width) | |
| path.points[:, 1] /= (h) # / canvas_height) | |
| path.points = 0.5 * (path.points + 1.0) * canvas_width | |
| center_x, center_y = canvas_width / 2, canvas_height / 2 | |
| path.points[:, 0] += (params["original_center_x"] * canvas_width - center_x) | |
| path.points[:, 1] += (params["original_center_y"] * canvas_height - center_y) | |
| if multiply: | |
| canvas_width *= 2 | |
| canvas_height *= 2 | |
| for path in shapes: | |
| path.points *= 2 | |
| path.stroke_width *= multiply | |
| _render = pydiffvg.RenderFunction.apply | |
| scene_args = pydiffvg.RenderFunction.serialize_scene( | |
| canvas_width, canvas_height, shapes, shape_groups) | |
| img = _render(canvas_width, # width | |
| canvas_height, # height | |
| 2, # num_samples_x | |
| 2, # num_samples_y | |
| 0, # seed | |
| None, | |
| *scene_args) | |
| img = img[:, :, 3:4] * img[:, :, :3] + \ | |
| torch.ones(img.shape[0], img.shape[1], 3, | |
| device=device) * (1 - img[:, :, 3:4]) | |
| img = img[:, :, :3].cpu().numpy() | |
| return img | |