Source / utilsV4.py
Richard1231's picture
Upload folder using huggingface_hub
470ac18
'''
Guowang Xie
from utilsV3.py
'''
import torch
from torch.utils import data
from torch.autograd import Variable, Function
import numpy as np
import sys, os, math
import cv2
import time
import re
import random
from scipy.interpolate import griddata
from tpsV2 import createThinPlateSplineShapeTransformer
def adjust_position(x_min, y_min, x_max, y_max, new_shape):
if (new_shape[0] - (x_max - x_min)) % 2 == 0:
f_g_0_0 = (new_shape[0] - (x_max - x_min)) // 2
f_g_0_1 = f_g_0_0
else:
f_g_0_0 = (new_shape[0] - (x_max - x_min)) // 2
f_g_0_1 = f_g_0_0 + 1
if (new_shape[1] - (y_max - y_min)) % 2 == 0:
f_g_1_0 = (new_shape[1] - (y_max - y_min)) // 2
f_g_1_1 = f_g_1_0
else:
f_g_1_0 = (new_shape[1] - (y_max - y_min)) // 2
f_g_1_1 = f_g_1_0 + 1
# return f_g_0_0, f_g_0_1, f_g_1_0, f_g_1_1
return f_g_0_0, f_g_1_0, new_shape[0] - f_g_0_1, new_shape[1] - f_g_1_1
def get_matric_edge(matric):
return np.concatenate((matric[:, 0, :], matric[:, -1, :], matric[0, 1:-1, :], matric[-1, 1:-1, :]), axis=0)
class SaveFlatImage(object):
'''
Post-processing and save result.
Function:
flatByRegressWithClassiy_multiProcessV2: Selecting a post-processing method
flatByfiducial_TPS: Thin Plate Spline, input multi-batch
flatByfiducial_interpolation: Interpolation, input one image
'''
def __init__(self, path, date, date_time, _re_date, data_path_validate, data_path_test, batch_size, preproccess=False, postprocess='tps_gpu', device=torch.device('cuda:0')):
self.path = path
self.date = date
self.date_time = date_time
self._re_date = _re_date
self.preproccess = preproccess
self.data_path_validate =data_path_validate
self.data_path_test = data_path_test
self.batch_size = batch_size
self.device = device
self.col_gap = 0 #4
self.row_gap = self.col_gap# col_gap + 1 if col_gap < 6 else col_gap
# fiducial_point_gaps = [1, 2, 3, 4, 5, 6, 10, 12, 15, 20, 30, 60] # POINTS NUM: 61, 31, 21, 16, 13, 11, 7, 6, 5, 4, 3, 2
self.fiducial_point_gaps = [1, 2, 3, 5, 6, 10, 15, 30] # POINTS NUM: 31, 16, 11, 7, 6, 4, 3, 2
self.fiducial_point_num = [31, 16, 11, 7, 6, 4, 3, 2]
self.fiducial_num = self.fiducial_point_num[self.col_gap], self.fiducial_point_num[self.row_gap]
map_shape = (320, 320)
self.postprocess = postprocess
if self.postprocess == 'tps':
self.tps = createThinPlateSplineShapeTransformer(map_shape, fiducial_num=self.fiducial_num, device=self.device)
def location_mark(self, img, location, color=(0, 0, 255)):
stepSize = 0
for l in location.astype(np.int64).reshape(-1, 2):
cv2.circle(img,
(l[0] + math.ceil(stepSize / 2), l[1] + math.ceil(stepSize / 2)), 3, color, -1)
return img
def flatByfiducial_TPS(self, fiducial_points, segment, im_name, epoch, perturbed_img=None, scheme='validate', is_scaling=False):
'''
flat_shap controls the output image resolution
'''
# if (scheme == 'test' or scheme == 'eval') and is_scaling:
# pass
# else:
if scheme == 'test' or scheme == 'eval':
perturbed_img_path = self.data_path_test + im_name
perturbed_img = cv2.imread(perturbed_img_path, flags=cv2.IMREAD_COLOR)
perturbed_img = cv2.resize(perturbed_img, (960, 1024))
elif scheme == 'validate' and perturbed_img is None:
RGB_name = im_name.replace('gw', 'png')
perturbed_img_path = self.data_path_validate + '/png/' + RGB_name
perturbed_img = cv2.imread(perturbed_img_path, flags=cv2.IMREAD_COLOR)
elif perturbed_img is not None:
perturbed_img = perturbed_img.transpose(1, 2, 0)
fiducial_points = fiducial_points / [992, 992]
perturbed_img_shape = perturbed_img.shape[:2]
sshape = fiducial_points[::self.fiducial_point_gaps[self.row_gap], ::self.fiducial_point_gaps[self.col_gap], :]
flat_shap = segment * [self.fiducial_point_gaps[self.col_gap], self.fiducial_point_gaps[self.row_gap]] * [self.fiducial_point_num[self.col_gap], self.fiducial_point_num[self.row_gap]]
# flat_shap = perturbed_img_shape
time_1 = time.time()
perturbed_img_ = torch.tensor(perturbed_img.transpose(2,0,1)[None,:])
fiducial_points_ = (torch.tensor(fiducial_points.transpose(1, 0,2).reshape(-1, 2))[None,:]-0.5)*2
rectified = self.tps(perturbed_img_.double().to(self.device), fiducial_points_.to(self.device), list(flat_shap))
time_2 = time.time()
time_interval = time_2 - time_1
print('TPS time: '+ str(time_interval))
flat_img = rectified[0].cpu().numpy().transpose(1,2,0)
'''save'''
flat_img = flat_img.astype(np.uint8)
i_path = os.path.join(self.path, self.date + self.date_time + ' @' + self._re_date,
str(epoch)) if self._re_date is not None else os.path.join(self.path,
self.date + self.date_time,
str(epoch))
''''''
perturbed_img_mark = self.location_mark(perturbed_img.copy(), sshape*perturbed_img_shape[::-1], (0, 0, 255))
if scheme == 'test':
i_path += '/test'
if not os.path.exists(i_path):
os.makedirs(i_path)
im_name = im_name.replace('gw', 'png')
cv2.imwrite(i_path + '/mark_' + im_name, perturbed_img_mark)
cv2.imwrite(i_path + '/' + im_name, flat_img)
def flatByfiducial_interpolation(self, fiducial_points, segment, im_name, epoch, perturbed_img=None, scheme='validate', is_scaling=False):
''''''
if scheme == 'test' or scheme == 'eval':
perturbed_img_path = self.data_path_test + im_name
perturbed_img = cv2.imread(perturbed_img_path, flags=cv2.IMREAD_COLOR)
perturbed_img = cv2.resize(perturbed_img, (960, 1024))
elif scheme == 'validate' and perturbed_img is None:
RGB_name = im_name.replace('gw', 'png')
perturbed_img_path = self.data_path_validate + '/png/' + RGB_name
perturbed_img = cv2.imread(perturbed_img_path, flags=cv2.IMREAD_COLOR)
elif perturbed_img is not None:
perturbed_img = perturbed_img.transpose(1, 2, 0)
fiducial_points = fiducial_points / [992, 992] * [960, 1024]
col_gap = 2 #4
row_gap = col_gap# col_gap + 1 if col_gap < 6 else col_gap
# fiducial_point_gaps = [1, 2, 3, 4, 5, 6, 10, 12, 15, 20, 30, 60] # POINTS NUM: 61, 31, 21, 16, 13, 11, 7, 6, 5, 4, 3, 2
fiducial_point_gaps = [1, 2, 3, 5, 6, 10, 15, 30] # POINTS NUM: 31, 16, 11, 7, 6, 4, 3, 2
sshape = fiducial_points[::fiducial_point_gaps[row_gap], ::fiducial_point_gaps[col_gap], :]
segment_h, segment_w = segment * [fiducial_point_gaps[col_gap], fiducial_point_gaps[row_gap]]
fiducial_points_row, fiducial_points_col = sshape.shape[:2]
im_x, im_y = np.mgrid[0:(fiducial_points_col - 1):complex(fiducial_points_col),
0:(fiducial_points_row - 1):complex(fiducial_points_row)]
tshape = np.stack((im_x, im_y), axis=2) * [segment_w, segment_h]
tshape = tshape.reshape(-1, 2)
sshape = sshape.reshape(-1, 2)
output_shape = (segment_h * (fiducial_points_col - 1), segment_w * (fiducial_points_row - 1))
grid_x, grid_y = np.mgrid[0:output_shape[0] - 1:complex(output_shape[0]),
0:output_shape[1] - 1:complex(output_shape[1])]
time_1 = time.time()
# grid_z = griddata(tshape, sshape, (grid_y, grid_x), method='cubic').astype('float32')
grid_ = griddata(tshape, sshape, (grid_y, grid_x), method='linear').astype('float32')
flat_img = cv2.remap(perturbed_img, grid_[:, :, 0], grid_[:, :, 1], cv2.INTER_CUBIC)
time_2 = time.time()
time_interval = time_2 - time_1
print('Interpolation time: '+ str(time_interval))
''''''
flat_img = flat_img.astype(np.uint8)
i_path = os.path.join(self.path, self.date + self.date_time + ' @' + self._re_date,
str(epoch)) if self._re_date is not None else os.path.join(self.path,
self.date + self.date_time,
str(epoch))
''''''
perturbed_img_mark = self.location_mark(perturbed_img.copy(), sshape, (0, 0, 255))
shrink_paddig = 0 # 2 * edge_padding
x_start, x_end, y_start, y_end = shrink_paddig, segment_h * (fiducial_points_col - 1) - shrink_paddig, shrink_paddig, segment_w * (fiducial_points_row - 1) - shrink_paddig
x_ = (perturbed_img_mark.shape[0]-(x_end-x_start))//2
y_ = (perturbed_img_mark.shape[1]-(y_end-y_start))//2
flat_img_new = np.zeros_like(perturbed_img_mark)
flat_img_new[x_:perturbed_img_mark.shape[0] - x_, y_:perturbed_img_mark.shape[1] - y_] = flat_img
img_figure = np.concatenate(
(perturbed_img_mark, flat_img_new), axis=1)
if scheme == 'test':
i_path += '/test'
if not os.path.exists(i_path):
os.makedirs(i_path)
im_name = im_name.replace('gw', 'png')
cv2.imwrite(i_path + '/' + im_name, img_figure)
def flatByRegressWithClassiy_multiProcessV2(self, pred_fiducial_points, pred_segment, im_name, epoch, process_pool=None, perturbed_img=None, scheme='validate', is_scaling=False):
for i_val_i in range(pred_fiducial_points.shape[0]):
if self.postprocess == 'tps':
self.flatByfiducial_TPS(pred_fiducial_points[i_val_i], pred_segment[i_val_i], im_name[i_val_i], epoch, None if perturbed_img is None else perturbed_img[i_val_i], scheme, is_scaling)
elif self.postprocess == 'interpolation':
self.flatByfiducial_interpolation(pred_fiducial_points[i_val_i], pred_segment[i_val_i], im_name[i_val_i], epoch, None if perturbed_img is None else perturbed_img[i_val_i], scheme, is_scaling)
else:
print('Error: Other postprocess.')
exit()
class AverageMeter(object):
"""Computes and stores the average and current value"""
def __init__(self):
self.reset()
def reset(self):
self.val = 0
self.avg = 0
self.sum = 0
self.count = 0
def update(self, val, n=1, m=1):
self.val = val
self.sum += val * m
self.count += n
self.avg = self.sum / self.count
class FlatImg(object):
'''
args:
self.save_flat_mage:Initialize the post-processing. Select a method in "postprocess_list".
'''
def __init__(self, args, path, date, date_time, _re_date, model,\
reslut_file, n_classes, optimizer, \
model_D=None, optimizer_D=None, \
loss_fn=None, loss_fn2=None, data_loader=None, data_loader_hdf5=None, dataPackage_loader = None, \
data_path=None, data_path_validate=None, data_path_test=None, data_preproccess=True): #, valloaderSet, v_loaderSet
self.args = args
self.path = path
self.date = date
self.date_time = date_time
self._re_date = _re_date
# self.valloaderSet = valloaderSet
# self.v_loaderSet = v_loaderSet
self.model = model
self.model_D = model_D
self.reslut_file = reslut_file
self.n_classes = n_classes
self.optimizer = optimizer
self.optimizer_D = optimizer_D
self.loss_fn = loss_fn
self.loss_fn2 = loss_fn2
self.data_loader = data_loader
self.data_loader_hdf5 = data_loader_hdf5
self.dataPackage_loader = dataPackage_loader
self.data_path = data_path
self.data_path_validate = data_path_validate
self.data_path_test = data_path_test
self.data_preproccess = data_preproccess
postprocess_list = ['tps', 'interpolation']
self.save_flat_mage = SaveFlatImage(self.path, self.date, self.date_time, self._re_date, self.data_path_validate, self.data_path_test, self.args.batch_size, self.data_preproccess, postprocess=postprocess_list[0], device=torch.device(self.args.device))
self.validate_loss = AverageMeter()
self.validate_loss_regress = AverageMeter()
self.validate_loss_segment = AverageMeter()
self.lambda_loss = 1
self.lambda_loss_segment = 1
self.lambda_loss_a = 1
self.lambda_loss_b = 1
self.lambda_loss_c = 1
def loadTrainData(self, data_split, is_shuffle=True):
train_loader = self.data_loader(self.data_path, split=data_split, img_shrink=self.args.img_shrink, preproccess=self.data_preproccess)
trainloader = data.DataLoader(train_loader, batch_size=self.args.batch_size, num_workers=min(self.args.batch_size, 8), drop_last=True, pin_memory=True,
shuffle=is_shuffle)
return trainloader
def loadValidateAndTestData(self, is_shuffle=True, sub_dir='shrink_512/crop/'):
v1_loader = self.data_loader(self.data_path_validate, split='validate', img_shrink=self.args.img_shrink, is_return_img_name=True, preproccess=self.data_preproccess)
valloader1 = data.DataLoader(v1_loader, batch_size=self.args.batch_size, num_workers=min(self.args.batch_size, 8), pin_memory=True, \
shuffle=is_shuffle)
'''val sets'''
v_loaderSet = {
'v1_loader': v1_loader,
}
valloaderSet = {
'valloader1': valloader1,
}
# sub_dir = 'crop/crop/'
t1_loader = self.data_loader(self.data_path_test, split='test', img_shrink=self.args.img_shrink, is_return_img_name=True)
testloader1 = data.DataLoader(t1_loader, batch_size=self.args.batch_size, num_workers=self.args.batch_size, pin_memory=True, \
shuffle=False)
'''test sets'''
t_loaderSet = {
't1_loader': v1_loader,
}
testloaderSet = {
'testloader1': testloader1,
}
self.valloaderSet = valloaderSet
self.v_loaderSet = v_loaderSet
self.testloaderSet = testloaderSet
self.t_loaderSet = t_loaderSet
# return v_loaderSet, valloaderSet
def loadTestData(self, is_shuffle=True):
t1_loader = self.data_loader(self.data_path_test, split='test', img_shrink=self.args.img_shrink,
is_return_img_name=True)
testloader1 = data.DataLoader(t1_loader, batch_size=self.args.batch_size, num_workers=self.args.batch_size,
shuffle=False)
'''test sets'''
testloaderSet = {
'testloader1': testloader1,
}
self.testloaderSet = testloaderSet
def evalData(self, is_shuffle=True, sub_dir='shrink_512/crop/'):
eval_loader = self.data_loader(self.data_path_test, split='eval', img_shrink=self.args.img_shrink, is_return_img_name=True)
evalloader = data.DataLoader(eval_loader, batch_size=self.args.batch_size, num_workers=self.args.batch_size, pin_memory=True, \
shuffle=False)
self.evalloaderSet = evalloader
# return v_loaderSet, valloaderSet
def saveModel_epoch(self, epoch):
epoch += 1
state = {'epoch': epoch,
'model_state': self.model.state_dict(),
'optimizer_state': self.optimizer.state_dict(), # AN ERROR HAS OCCURED
}
i_path = os.path.join(self.path, self.date + self.date_time + ' @' + self._re_date,
str(epoch)) if self._re_date is not None else os.path.join(self.path, self.date + self.date_time, str(epoch))
if not os.path.exists(i_path):
os.makedirs(i_path)
if self._re_date is None:
torch.save(state, i_path + '/' + self.date + self.date_time + "{}".format(self.args.arch) + ".pkl") # "./trained_model/{}_{}_best_model.pkl"
else:
torch.save(state,
i_path + '/' + self._re_date + "@" + self.date + self.date_time + "{}".format(
self.args.arch) + ".pkl")
def validateOrTestModelV3(self, epoch, trian_t, validate_test='v_l2', is_scaling=False):
if validate_test == 'v_l4':
loss_segment_list = 0
loss_overall_list = 0
loss_local_list = 0
loss_edge_list = 0
loss_rectangles_list = 0
loss_list = []
begin_test = time.time()
with torch.no_grad():
for i_valloader, valloader in enumerate(self.valloaderSet.values()):
for i_val, (images, labels, segment, im_name) in enumerate(valloader):
try:
# save_img_ = random.choices([True, False], weights=[1, 0])[0]
save_img_ = random.choices([True, False], weights=[0.05, 0.95])[0]
# save_img_ = True
images = Variable(images)
labels = Variable(labels.cuda(self.args.device))
segment = Variable(segment.cuda(self.args.device))
outputs, outputs_segment = self.model(images)
loss_overall, loss_local, loss_edge, loss_rectangles = self.loss_fn(outputs, labels, size_average=True)
loss_segment = self.loss_fn2(outputs_segment, segment)
loss = self.lambda_loss * (loss_overall + loss_local + loss_edge * self.lambda_loss_a + loss_rectangles * self.lambda_loss_b) + self.lambda_loss_segment * loss_segment
# loss = self.lambda_loss * (loss_local + loss_rectangles + loss_edge*self.lambda_loss_a + loss_overall*self.lambda_loss_b) + self.lambda_loss_segment * loss_segment
pred_regress = outputs.data.cpu().numpy().transpose(0, 2, 3, 1) # (4, 1280, 1024, 2)
pred_segment = outputs_segment.data.round().int().cpu().numpy() # (4, 1280, 1024) ==outputs.data.argmax(dim=0).cpu().numpy()
if save_img_:
self.save_flat_mage.flatByRegressWithClassiy_multiProcessV2(pred_regress,
pred_segment, im_name,
epoch + 1,
perturbed_img=images.numpy(), scheme='validate', is_scaling=is_scaling)
loss_list.append(loss.item())
loss_segment_list += loss_segment.item()
loss_overall_list += loss_overall.item()
loss_local_list += loss_local.item()
# loss_edge_list += loss_edge.item()
# loss_rectangles_list += loss_rectangles.item()
except:
print('* save image validated error :'+im_name[0])
test_time = time.time() - begin_test
# if always_save_model:
# self.saveModel(epoch, save_path=self.path)
list_len = len(loss_list)
print('train time : {trian_t:.3f}\t'
'validate time : {test_time:.3f}\t'
'[o:{overall_avg:.4f} l:{local_avg:.4f} e:{edge_avg:.4f} r:{rectangles_avg:.4f}\t'
'[{loss_regress:.4f} {loss_segment:.4f}]\n'.format(
trian_t=trian_t, test_time=test_time,
overall_avg=loss_overall_list / list_len, local_avg=loss_local_list / list_len, edge_avg=loss_edge_list / list_len, rectangles_avg=loss_rectangles_list / list_len,
loss_regress=(loss_overall_list+loss_local_list+loss_edge_list) / list_len, loss_segment=loss_segment_list / list_len))
print('train time : {trian_t:.3f}\t'
'validate time : {test_time:.3f}\t'
'[o:{overall_avg:.4f} l:{local_avg:.4f} e:{edge_avg:.4f} r:{rectangles_avg:.4f}\t'
'[{loss_regress:.4f} {loss_segment:.4f}]\n'.format(
trian_t=trian_t, test_time=test_time,
overall_avg=loss_overall_list / list_len, local_avg=loss_local_list / list_len, edge_avg=loss_edge_list / list_len, rectangles_avg=loss_rectangles_list / list_len,
loss_regress=(loss_overall_list+loss_local_list+loss_edge_list) / list_len, loss_segment=loss_segment_list / list_len), file=self.reslut_file)
elif validate_test == 't_all':
begin_test = time.time()
with torch.no_grad():
for i_valloader, valloader in enumerate(self.testloaderSet.values()):
for i_val, (images, im_name) in enumerate(valloader):
try:
# save_img_ = True
save_img_ = random.choices([True, False], weights=[1, 0])[0]
# save_img_ = random.choices([True, False], weights=[0.2, 0.8])[0]
if save_img_:
images = Variable(images)
outputs, outputs_segment = self.model(images)
# outputs, outputs_segment = self.model(images, is_softmax=True)
pred_regress = outputs.data.cpu().numpy().transpose(0, 2, 3, 1)
pred_segment = outputs_segment.data.round().int().cpu().numpy() # (4, 1280, 1024) ==outputs.data.argmax(dim=0).cpu().numpy()
self.save_flat_mage.flatByRegressWithClassiy_multiProcessV2(pred_regress,
pred_segment, im_name,
epoch + 1,
scheme='test', is_scaling=is_scaling)
except:
print('* save image tested error :' + im_name[0])
test_time = time.time() - begin_test
print('test time : {test_time:.3f}'.format(
test_time=test_time))
print('test time : {test_time:.3f}'.format(
test_time=test_time),
file=self.reslut_file)
else:
begin_test = time.time()
with torch.no_grad():
for i_valloader, valloader in enumerate(self.testloaderSet.values()):
for i_val, (images, im_name) in enumerate(valloader):
try:
# save_img_ = True
# save_img_ = random.choices([True, False], weights=[1, 0])[0]
save_img_ = random.choices([True, False], weights=[0.4, 0.6])[0]
if save_img_:
images = Variable(images)
outputs, outputs_segment = self.model(images)
# outputs, outputs_segment = self.model(images, is_softmax=True)
pred_regress = outputs.data.cpu().numpy().transpose(0, 2, 3, 1)
pred_segment = outputs_segment.data.round().int().cpu().numpy() # (4, 1280, 1024) ==outputs.data.argmax(dim=0).cpu().numpy()
self.save_flat_mage.flatByRegressWithClassiy_multiProcessV2(pred_regress,
pred_segment, im_name,
epoch + 1,
scheme='test', is_scaling=is_scaling)
except:
print('* save image tested error :' + im_name[0])
test_time = time.time() - begin_test
print('test time : {test_time:.3f}'.format(
test_time=test_time))
print('test time : {test_time:.3f}'.format(
test_time=test_time),
file=self.reslut_file)