From 9db40d23f4a4f09268859277e7b12e80f20f53e7 Mon Sep 17 00:00:00 2001 From: Nikolay Date: Fri, 14 May 2021 05:48:35 +0300 Subject: [PATCH] PEP8 style fixes. Both cpu and gpu are now supported: changed from .cuda() to .to(device). Unused variables are removed. Path to TB logs is moved to args instead of hard-code. One TODO is added to fix the unimplemented function. --- cleaning/models/SmallUnet/unet.py | 33 ++++++------- cleaning/scripts/fine_tuning.py | 17 ++++--- .../fine_tuning_two_network_added_part.py | 38 ++++++--------- cleaning/scripts/run.py | 3 +- merging/merging_for_curves.py | 3 +- merging/merging_for_lines.py | 3 +- merging/utils/merging_for_lines_functions.py | 47 ++++++++++--------- 7 files changed, 70 insertions(+), 74 deletions(-) diff --git a/cleaning/models/SmallUnet/unet.py b/cleaning/models/SmallUnet/unet.py index d5b6db6..3f69ad7 100644 --- a/cleaning/models/SmallUnet/unet.py +++ b/cleaning/models/SmallUnet/unet.py @@ -4,13 +4,14 @@ import gc import torch -import torch, torch.nn as nn +import torch +import torch.nn as nn import torch.functional as F from torch.autograd import Variable import numpy as np -def CreateConvBnRelu(in_channels, out_channels, dilation=1): +def create_conv_bn_relu(in_channels, out_channels, dilation=1): module = nn.Sequential() module.add_module('conv', nn.Conv2d( in_channels=in_channels, @@ -30,35 +31,35 @@ def __init__(self): super(SmallUnet, self).__init__() # encoder - self.add_module('EncConvBnRelu1_1', CreateConvBnRelu(3, 64)) - self.add_module('EncConvBnRelu1_2', CreateConvBnRelu(64, 64)) + self.add_module('EncConvBnRelu1_1', create_conv_bn_relu(3, 64)) + self.add_module('EncConvBnRelu1_2', create_conv_bn_relu(64, 64)) self.add_module('EncMp1', nn.MaxPool2d(kernel_size=2)) - self.add_module('EncConvBnRelu2_1', CreateConvBnRelu(64, 128)) - self.add_module('EncConvBnRelu2_2', CreateConvBnRelu(128, 128)) + self.add_module('EncConvBnRelu2_1', create_conv_bn_relu(64, 128)) + self.add_module('EncConvBnRelu2_2', create_conv_bn_relu(128, 128)) self.add_module('EncMp2', nn.MaxPool2d(kernel_size=2)) - self.add_module('EncConvBnRelu3_1', CreateConvBnRelu(128, 256)) - self.add_module('EncConvBnRelu3_2', CreateConvBnRelu(256, 256)) + self.add_module('EncConvBnRelu3_1', create_conv_bn_relu(128, 256)) + self.add_module('EncConvBnRelu3_2', create_conv_bn_relu(256, 256)) self.add_module('EncMp3', nn.MaxPool2d(kernel_size=2)) # lowest layer - self.add_module('ConvBnRelu4_1', CreateConvBnRelu(256, 512)) - self.add_module('ConvBnRelu4_2', CreateConvBnRelu(512, 512)) + self.add_module('ConvBnRelu4_1', create_conv_bn_relu(256, 512)) + self.add_module('ConvBnRelu4_2', create_conv_bn_relu(512, 512)) self.add_module('Us4', nn.Upsample(scale_factor=2)) # decoder - self.add_module('DecConvBnRelu3_1', CreateConvBnRelu(512 + 256, 256)) - self.add_module('DecConvBnRelu3_2', CreateConvBnRelu(256, 256)) + self.add_module('DecConvBnRelu3_1', create_conv_bn_relu(512 + 256, 256)) + self.add_module('DecConvBnRelu3_2', create_conv_bn_relu(256, 256)) self.add_module('DecUs3', nn.Upsample(scale_factor=2)) - self.add_module('DecConvBnRelu2_1', CreateConvBnRelu(256 + 128, 128)) - self.add_module('DecConvBnRelu2_2', CreateConvBnRelu(128, 128)) + self.add_module('DecConvBnRelu2_1', create_conv_bn_relu(256 + 128, 128)) + self.add_module('DecConvBnRelu2_2', create_conv_bn_relu(128, 128)) self.add_module('DecUs2', nn.Upsample(scale_factor=2)) # prediction - self.add_module('PredConvBnRelu_1', CreateConvBnRelu(128 + 64, 64)) - self.add_module('PredConvBnRelu_2', CreateConvBnRelu(64, 64)) + self.add_module('PredConvBnRelu_1', create_conv_bn_relu(128 + 64, 64)) + self.add_module('PredConvBnRelu_2', create_conv_bn_relu(64, 64)) self.add_module('PredDense', nn.Conv2d( in_channels=64, out_channels=1, diff --git a/cleaning/scripts/fine_tuning.py b/cleaning/scripts/fine_tuning.py index dfd23e8..8c30534 100644 --- a/cleaning/scripts/fine_tuning.py +++ b/cleaning/scripts/fine_tuning.py @@ -20,7 +20,7 @@ MODEL_LOSS = { 'UNET': { - 'model': UNet(n_channels=3,n_classes=1,final_tanh=True), + 'model': UNet(n_channels=3, n_classes=1, final_tanh=True), 'loss': CleaningLoss(kind='BCE', with_restore=False) }, 'SmallUNET': { @@ -28,7 +28,7 @@ 'loss': CleaningLoss(kind='BCE', with_restore=False) }, 'UNET_MSE': { - 'model': UNet(n_channels=3,n_classes=1), + 'model': UNet(n_channels=3, n_classes=1), 'loss': CleaningLoss(kind='MSE', with_restore=False) }, @@ -38,23 +38,26 @@ def parse_args(): parser = argparse.ArgumentParser() - parser.add_argument('--model', type=str, required=True, help='What model to use, one of [ "UNET,"UNET_MSE,"SmallUnet"]', default='UNET') + parser.add_argument('--model', type=str, required=True, + help='What model to use, one of [ "UNET,"UNET_MSE,"SmallUnet"]', default='UNET') parser.add_argument('--n_epochs', type=int, help='Num of epochs for training', default=10) parser.add_argument('--datadir', type=str, help='Path to training dataset') parser.add_argument('--valdatadir', type=str, help='Path to validation dataset') parser.add_argument('--batch_size', type=int, default=8) parser.add_argument('--name', type=str, help='Name of the experiment') parser.add_argument('--model_path', type=str, default=None, help='Path to model checkpoint') + parser.add_argument('--tb_dir_prefix', type=str, + default='/logs/tb_logs_article/fine_tuning_', help='Path to TB dir and prefix of name') + + return parser.parse_args() - args = parser.parse_args() - return args def get_dataloaders(args): train_transform = transforms.Compose([ transforms.ToTensor(), transforms.ColorJitter() ]) - #TODO Make sure that this should be MakeDataSynt and not MakeData from dataloader.py + # TODO Make sure that this should be MakeDataSynt and not MakeData from dataloader.py dset_synt = MakeDataSynt(args.datadir, args.datadir, train_transform, 1) dset_val_synt = MakeDataSynt(args.valdatadir, args.valdatadir, train_transform) @@ -121,7 +124,7 @@ def load_model(model, path): def main(args): - tb_dir = '/logs/tb_logs_article/fine_tuning_' + args.name + tb_dir = args.tb_dir_prefix + args.name tb = SummaryWriter(tb_dir) train_loader, val_loader = get_dataloaders(args) diff --git a/cleaning/scripts/fine_tuning_two_network_added_part.py b/cleaning/scripts/fine_tuning_two_network_added_part.py index 87239ca..44dedbb 100644 --- a/cleaning/scripts/fine_tuning_two_network_added_part.py +++ b/cleaning/scripts/fine_tuning_two_network_added_part.py @@ -100,18 +100,15 @@ def validate(tb, val_loader, unet, gen, loss_func, global_step ): # Cleaning prediction logits_restor, logits_extract = None, unet(x_input) - #1 - Cleaning prediction + # 1 - Cleaning prediction logits_extract = 1 - logits_extract.unsqueeze(1) y_restor = 1. - y_restor.type(torch.FloatTensor).cuda().unsqueeze(1) - #generator prediction based on cleaning prediction + # generator prediction based on cleaning prediction logits_restore = gen.forward(logits_extract).unsqueeze(1) # restoration + extraction - - loss = loss_func(1 - (logits_extract + logits_restore), None, 1- y_restor,None ) - val_loss_epoch.append(loss.cpu().data.numpy()) iou_scr_without_gan = iou_score(1 - torch.round(logits_extract.squeeze(1)).cpu().long().numpy(),1 - torch.round(y_restor.squeeze(1)).cpu().long().numpy()) @@ -125,7 +122,6 @@ def validate(tb, val_loader, unet, gen, loss_func, global_step ): tb.add_scalar('val_iou_extract', np.mean(val_iou_extract), global_step=global_step) tb.add_scalar('val_loss', np.mean(val_loss_epoch), global_step=global_step) tb.add_scalar('val_iou_without_gan', np.mean(val_iou_without_gan), global_step=global_step) - out_grid = torchvision.utils.make_grid(1.- torch.clamp(logits_extract + logits_restore, 0, 1).cpu()) input_grid = torchvision.utils.make_grid(1. - logits_extract.cpu()) @@ -133,7 +129,6 @@ def validate(tb, val_loader, unet, gen, loss_func, global_step ): input_clean_grid = torchvision.utils.make_grid(x_input.cpu()) tb.add_image(tag='val_first_input', img_tensor=input_clean_grid, global_step=global_step) - tb.add_image(tag='val_out_extract', img_tensor=out_grid, global_step=global_step) tb.add_image(tag='val_input', img_tensor=input_grid, global_step=global_step) tb.add_image(tag='val_true', img_tensor=true_grid, global_step=global_step) @@ -156,16 +151,16 @@ def main(args): tb = SummaryWriter(tb_dir) train_loader, val_loader = get_dataloaders(args) - + device = 'cuda' if torch.cuda.is_available() else 'cpu' if args.model not in ["UNET"]: raise Exception('Unsupported type of model, choose from [ "UNET"]') gen = MODEL_LOSS[args.model]['gen'] - unet = MODEL_LOSS[args.model]['unet'] + unet = MODEL_LOSS[args.model]['unet'] loss_func = MODEL_LOSS[args.model]['loss'] - gen = gen.cuda() - unet = unet.cuda() + gen = gen.to(device) + unet = unet.to(device) if 'gen_path' in args and args.gen_path is not None: gen = load_model(gen, args.gen_path) @@ -187,44 +182,40 @@ def main(args): for x_input, y_extract, y_restor in tqdm(train_loader): # data reading # unet.train() - x_input = torch.FloatTensor(x_input).cuda() - y_extract = y_extract.type(torch.FloatTensor).cuda().unsqueeze(1) - y_restor = 1. - y_restor.type(torch.FloatTensor).cuda().unsqueeze(1) + x_input = torch.FloatTensor(x_input).to(device) + y_extract = y_extract.type(torch.FloatTensor).to(device).unsqueeze(1) + y_restor = 1. - y_restor.type(torch.FloatTensor).to(device).unsqueeze(1) unet.eval() with torch.no_grad(): - logits_restor, logits_extract = None, unet(x_input) + logits_extract = unet(x_input) logits_extract = 1 - logits_extract.unsqueeze(1) - logits_restore = gen.forward(logits_extract).unsqueeze(1) # restoration + extraction - # if Cleaning loss use this - gen_loss = loss_func(1 - (logits_extract + logits_restore), None, 1- y_restor,None ) + gen_loss = loss_func(1 - (logits_extract + logits_restore), None, 1-y_restor, None ) # else if with_restore =True use this # input_fake = torch.cat((logits_extract + logits_restore,logits_extract),dim = 1) # # gen_loss = loss_func(logits_extract, input_fake, y_extract, y_restor) - gen_opt.zero_grad() gen_loss.backward() gen_opt.step() gen_step += 1 - if(np.random.random() <=0.5): - disc_step+=1 + if np.random.random() <= 0.5: + disc_step += 1 global_step += 1 if global_step <= 1: continue - tb.add_scalar('gen_vectran_loss',gen_loss.item(), global_step=global_step) + tb.add_scalar('gen_vectran_loss', gen_loss.item(), global_step=global_step) # tb.add_scalar('train_loss', loss.cpu().data.numpy(), global_step=global_step) - if global_step % 100 == 0 or global_step <= 2: out_grid = torchvision.utils.make_grid(1. - torch.clamp(logits_extract + logits_restore, 0, 1).cpu()) input_grid = torchvision.utils.make_grid(1. - logits_extract.cpu()) @@ -248,7 +239,6 @@ def main(args): save_model(unet, os.path.join(tb_dir, 'unet_it_%s.pth' % global_step)) - if __name__ == '__main__': args = parse_args() main(args) diff --git a/cleaning/scripts/run.py b/cleaning/scripts/run.py index 9ea71a1..77491fe 100644 --- a/cleaning/scripts/run.py +++ b/cleaning/scripts/run.py @@ -112,7 +112,8 @@ def main(options): skio.imsave(options.cleaned_filename, cleaned_rgb) if options.vectorize: - vector_model = load_vector_model(options.vector_model_filename) + # TODO: implement of import load_vector_model() func + vector_model = load_vector_model(options.vector_model_filename) if options.use_patches: patches_rgb, patches_offsets = split_to_patches(cleaned_rgb, options.patch_size) #patches_vector = [] diff --git a/merging/merging_for_curves.py b/merging/merging_for_curves.py index 30785ab..a448984 100644 --- a/merging/merging_for_curves.py +++ b/merging/merging_for_curves.py @@ -5,13 +5,12 @@ import numpy as np -sys.path.append('/code') from util_files.job_tuples.calculate_results_for_curves import job_tuples from util_files.simplification.join_qb import join_quad_beziers from util_files.optimization.optimizer.logging import Logger from util_files.rendering.cairo import PT_LINE, PT_QBEZIER from util_files.data.graphics.graphics import Path, VectorImage - +sys.path.append('/code') def main(options, width_percentile=90, fit_tol=.5, w_tol=np.inf, join_tol=.5): logger = Logger.prepare_logger(loglevel='info', logfile=None) diff --git a/merging/merging_for_lines.py b/merging/merging_for_lines.py index b625cf2..8d1bea1 100644 --- a/merging/merging_for_lines.py +++ b/merging/merging_for_lines.py @@ -4,6 +4,7 @@ import argparse import torch + def postprocess(y_pred_render, patches_offsets, input_rgb, cleaned_image, it, options): ''' @@ -38,6 +39,4 @@ def postprocess(y_pred_render, patches_offsets, input_rgb, cleaned_image, it, op save_svg(result_tuning, cleaned_image.shape, options.image_name[it], options.output_dir + 'iou_postprocess/') result_tuning = lines_matching(result_tuning, frac=0.07) save_svg(result_tuning, cleaned_image.shape, options.image_name[it], options.output_dir + 'lines_matching/') - - return result_tuning diff --git a/merging/utils/merging_for_lines_functions.py b/merging/utils/merging_for_lines_functions.py index 5f6a58a..aa65fb2 100644 --- a/merging/utils/merging_for_lines_functions.py +++ b/merging/utils/merging_for_lines_functions.py @@ -1,18 +1,11 @@ #!/usr/bin/env python3 from __future__ import division - -import os, sys - -sys.path.append("..") -sys.path.append(os.path.join(os.getcwd(), '..')) - from tqdm import tqdm - +import os +import sys import numpy as np import os import math - - from scipy.spatial import distance from sklearn.linear_model import LinearRegression @@ -20,10 +13,10 @@ import util_files.data.graphics_primitives as graphics_primitives from util_files.rendering.cairo import render_with_skeleton - from util_files.geometric import liang_barsky_screen - +sys.path.append("..") +sys.path.append(os.path.join(os.getcwd(), '..')) def ordered(line): @@ -34,6 +27,7 @@ def ordered(line): return np.array([min_x, min_y, max_x, max_y]) + def clip_to_box(y_pred, box_size=(64, 64)): width, height = box_size bbox = (0, 0, width, height) @@ -42,7 +36,7 @@ def clip_to_box(y_pred, box_size=(64, 64)): clipped_point1, clipped_point2, is_drawn = \ liang_barsky_screen(point1, point2, bbox) except: - np.asarray([np.nan, np.nan, np.nan, np.nan, np.nan, np.nan]) + return np.asarray([np.nan, np.nan, np.nan, np.nan, np.nan, np.nan]) if (clipped_point1 and clipped_point2): return np.asarray([clipped_point1, clipped_point2, y_pred[4:]]).ravel() @@ -63,6 +57,7 @@ def assemble_vector_patches(patches_vector, patches_offsets): primitives.append(patch_vector) return np.array(primitives) + def tensor_vector_graph_numpy(y_pred_render, patches_offsets, options): nump = np.array(list(map(clip_to_box, y_pred_render.reshape(-1, 6).cpu().detach().numpy()))) nump = assemble_vector_patches(np.array((nump.reshape(-1, options.model_output_count, 6))), @@ -73,9 +68,9 @@ def tensor_vector_graph_numpy(y_pred_render, patches_offsets, options): nump = nump[(nump[:, -2] > 0.3)] nump = nump[(nump[:, -1] > 0.5)] - print('nump',nump.shape) + print('nump', nump.shape) nump = nump[((nump[:, 0] - nump[:, 2]) ** 2 + (nump[:, 1] - nump[:, 3]) ** 2 >= 3)] - print('nump',nump.shape) + print('nump', nump.shape) return nump @@ -125,6 +120,7 @@ def merge_close_lines(lines, threshold=0.5): return np.array([dt[0], y_pred[0], dt[-1], y_pred[-1]]) + def point_to_line_distance(point, line): px, py = point x1, y1, x2, y2 = line @@ -134,6 +130,7 @@ def point_to_line_distance(point, line): return math.hypot(px - x1, py - y1) return np.abs(dy * px - dx * py + x2 * y1 - x1 * y2) / np.sqrt(dy ** 2 + dx ** 2) + def point_segment_distance(point, line): px, py = point x1, y1, x2, y2 = line @@ -161,6 +158,7 @@ def point_segment_distance(point, line): return math.hypot(dx, dy) + def dist(line0, line1): if (point_to_line_distance(line0[:2], line1[:4]) >= 2 or point_to_line_distance(line0[2:4], line1[ :4]) >= 2 or point_to_line_distance( @@ -187,6 +185,7 @@ def dfs(graph, start): def line_legth(line): return np.sqrt((line[0] - line[2]) ** 2 + (line[1] - line[3]) ** 2) + def intersect(line0, line1): # Solve the system [p1-p0, q1-q0]*[t1, t2]^T = q0 - p0 # where line0 = (p0, p1) and line1 = (q0, q1) @@ -206,6 +205,7 @@ def intersect(line0, line1): return [(t1, t2)] return [] + def angle_radians(pt1, pt2): x1, y1 = pt1 x2, y2 = pt2 @@ -218,6 +218,7 @@ def angle_radians(pt1, pt2): def normalize(x): return x / np.linalg.norm(x) + def compute_angle(line0, line1): pt1 = normalize([line0[2] - line0[0], line0[3] - line0[1]]) pt2 = normalize([line1[2] - line1[0], line1[3] - line1[1]]) @@ -226,12 +227,13 @@ def compute_angle(line0, line1): angle = math.degrees(angle_radians(pt1, pt2)) except: angle = 0 - if (angle >= 90 and angle <= 270): + if angle >= 90 and angle <= 270: angle = np.abs(180 - angle) - elif (angle > 270 and angle <= 360): + elif angle > 270 and angle <= 360: angle = 360 - angle return angle + def merge_close(lines, idx, widths, tol=1e-3, max_dist=5, max_angle=15, window_width=100): window = [-window_width, -window_width, window_width, window_width] n = len(lines) @@ -255,7 +257,7 @@ def merge_close(lines, idx, widths, tol=1e-3, max_dist=5, max_angle=15, window_w for i in range(n): - if (line_legth(lines[i, :4]) < 3): + if line_legth(lines[i, :4]) < 3: continue elif (close[i]) and (i not in merged): @@ -268,6 +270,7 @@ def merge_close(lines, idx, widths, tol=1e-3, max_dist=5, max_angle=15, window_w result.append((lines[i])) return result + def draw_with_skeleton(lines, drawing_scale=1, skeleton_line_width=0, skeleton_node_size=0, max_x=64, max_y=64): scaled_primitives = lines.copy() scaled_primitives[..., :-1] *= drawing_scale @@ -296,11 +299,14 @@ def maximiz_final_iou(nump, input_rgb): mse_ref = tmp_scr l += 1 return lines -def two_point_dist(p1,p2): + + +def two_point_dist(p1, p2): p1 = np.array(p1) p2 = np.array(p2) return np.sqrt((np.square(p1-p2)).sum()) + def line(p1): A = (p1[1] - p1[3]) B = (p1[2] - p1[0]) @@ -318,6 +324,7 @@ def intersection(L1, L2): else: return [] + def lines_matching(lines, frac = 0.01): """ lines: lines @@ -346,12 +353,8 @@ def lines_matching(lines, frac = 0.01): return lines - - - def save_svg(result_vector, size, name, output_dir): print(size) os.makedirs(output_dir, exist_ok=True) img = draw_with_skeleton(result_vector, drawing_scale=1, skeleton_line_width=0, skeleton_node_size=0, max_x=64, max_y=64) return img -