def train(
model,
image,
D,
opt,
args,
):
t = trange(args.epochs)
B = args.batch_size
ls = args.latent_size
sz = args.size
points = torch.stack(torch.meshgrid(
torch.linspace(-1, 1, steps=args.size, device=device),
torch.linspace(-1, 1, steps=args.size, device=device),
), dim=-1)[None, ...].expand(B, -1, -1, -1)
for i in t:
opt.zero_grad()
# make multiple times and train them at the same time
time = torch.randn(B,1,1,1,device=device).expand(B,sz,sz,1)
latent = torch.randn(ls, device=device)[None, None, :]\
.expand([*points.shape[:-1],-1])
got = model(points, time, latent)
style_loss, content_loss = 0, 0
for g in got.split(1, dim=0):
sl, cl = D(g.permute(0, 3, 1, 2))
style_loss = style_loss + sl
#content_loss = content_loss + cl
loss = style_loss + content_loss
t.set_postfix(style=f"{style_loss:.03f}", content=f"{content_loss:.03f}")
loss.backward()
opt.step()
if i % args.valid_freq == 0:
save_image(f"outputs/fieldgan_{i:05}.png", got[0].clamp(min=0, max=1))
def test(model, args):
with torch.no_grad():
ls = args.latent_size
points = torch.stack(torch.meshgrid(
torch.linspace(-1, 1, steps=args.size, device=device),
torch.linspace(-1, 1, steps=args.size, device=device),
), dim=-1)
latent1 = torch.randn(ls, device=device)[None, None, :]\
.expand([*points.shape[:-1],-1])
latent2 = torch.randn(ls, device=device)[None, None, :]\
.expand([*points.shape[:-1],-1])
steps = 100
for i, t in enumerate(torch.linspace(0, 1, steps=steps, device=device)):
t = t.expand([*points.shape[:-1], 1])
alpha = i/steps
latent = (1-alpha) * latent1 + alpha * latent2
got = model(points, t, latent)
save_image(f"outputs/fieldgan_test_{i:03}.png", got)
latent3 = torch.randn(ls, device=device)[None, None, :]\
.expand([*points.shape[:-1],-1])
for i, t in enumerate(reversed(torch.linspace(0, 1, steps=steps, device=device))):
t = t.expand([*points.shape[:-1], 1])
alpha = i/steps
latent = (1-alpha) * latent2 + alpha * latent3
got = model(points, t, latent)
i = i + steps
save_image(f"outputs/fieldgan_test_{i:03}.png", got)
def main():
args = arguments()
with torch.no_grad():
model = torch.load(args.refl_model)
assert(hasattr(model, "refl")), "The provided model must have a refl"
r = model.refl
if isinstance(r, refl.LightAndRefl): r = r.refl
# just check the first choice for now, can add a flag for it later
if isinstance(r, refl.WeightedChoice): r = r.choices[args.weighted_refl_idx]
# check again if it's another lit item
if isinstance(r, refl.LightAndRefl): r = r.refl
assert(isinstance(r, refl.Rusin)), f"must provide a rusin refl, got {type(r)}"
degs = torch.stack(torch.meshgrid(
# theta_h
torch.linspace(0, 90, 256, device=device, dtype=torch.float),
# theta_d
torch.linspace(0, 90, 256, device=device, dtype=torch.float),
# phi_d
torch.linspace(0, 360, 256, device=device, dtype=torch.float),
), dim=-1)
rads = torch.deg2rad(degs)
for i, theta_h in enumerate(rads.split(1, dim=2)):
latent = torch.randn(*rads.shape[:-2], r.latent_size, device=device)
theta_h = theta_h.squeeze(2)
out = r.raw(theta_h, latent)
save_image(f"outputs/rusin_eval_{i:03}.png", out)
return
def get_output_image(self, content_path, style_path, options: dict):
start_time = time.time()
style_content_model = VGG19Model(CONTENT_LAYERS, STYLE_LAYERS)
content_image, style_image = [
load_image(path) for path in (content_path, style_path)
]
image = tf.Variable(get_white_noise_image(tf.shape(content_image)[1:])) \
if options['white_noise_input'] else tf.Variable(content_image)
style_targets = style_content_model(style_image)['style_outputs']
content_targets = style_content_model(content_image)['content_outputs']
opt = tf.keras.optimizers.Adam(learning_rate=options['learning_rate'],
beta_1=0.99,
epsilon=1e-1)
style_content_model.compile(opt)
for epoch in range(options['epochs']):
for step in range(options['steps']):
style_content_model.fit(
image,
content_targets=content_targets,
style_targets=style_targets,
content_layer_weights=[1],
style_layer_weights=options['style_layer_weights'],
content_weight=options['content_weight'],
style_weight=options['style_weight'],
variation_weight=options['variation_weight'])
self.update_state(state='PROGRESS',
meta={
'current':
options['steps'] * epoch + step,
'total':
options['steps'] * options['epochs'],
'elapsed_time':
"{:.1f}s".format(time.time() - start_time)
})
output_path = Path('./static/output') / (str(self.request.id) + '.png')
save_image(image, Path(output_path))
return {
'output_path': str(output_path),
'total_time': "{:.1f}s".format(time.time() - start_time)
}
def main(): args = arguments() model = SmoothImageApprox(latent_size=args.latent_size).to(device) image = load_image(args.image, [args.size, args.size]) save_image(f"outputs/fieldgan_ref.png", image) image = image.permute(2, 0, 1) # use same image for content and loss. D = StyleTransfer(image[None, ...], image[None, ...]).to(device) image = image.to(device) init_image = model.init_zero(image.permute(1,2,0)) save_image(f"outputs/fieldgan_init.png", init_image) opt = optim.Adam(model.displacement.parameters(), lr=1e-3, weight_decay=0) #opt = optim.Adam(model.parameters(), lr=1e-3, weight_decay=0) train(model, image, D, opt, args) # TODO render the image after, fixing a latent noise and iterating through time test(model, args)
def main():
args = arguments()
sz = args.size
with torch.no_grad():
model = torch.load(args.model)
exit()
assert (isinstance(model, RigNeRF)), "Can only project pts of RigNeRF"
labels, cam, _ = loaders.load(args, training=False, device=device)
for i in trange(labels.shape[0]):
c = cam[i:i + 1]
pt2d = c.project_pts(model.points, sz)
out = torch.zeros(sz, sz, 3, dtype=torch.float, device=device)
pixels = pt2d.long()
pixels = pixels[((0 <= pixels) & (pixels < sz)).all(dim=-1)]
out[pixels[:, 0], pixels[:, 1]] = 1
save_image(f"outputs/proj_pts_{i:03}.png", out)
return
def calibrateCameraFromImages(self, path):
print("Calibrating:")
objpoints = [] # 3d points in real world space
imgpoints = [] # 2d points in image plane.
images = glob.glob(path + "/*")
# setup toolbar
toolbar_width = len(images)
sys.stdout.write("[%s]" % (" " * toolbar_width))
sys.stdout.flush()
sys.stdout.write("\b" * (toolbar_width + 1)) # return to start of line, after '['
# Step through the list and search for chessboard corners
for idx, file_name in enumerate(images):
img = cv2.imread(file_name)
gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
# Find the chessboard corners
ret, corners = cv2.findChessboardCorners(gray, self.chessboardSize, None)
# If found, add object points, image points
if ret:
objpoints.append(self.objp)
imgpoints.append(corners) # Draw and display the corners
cv2.drawChessboardCorners(img, self.chessboardSize, corners, ret)
# save image
save_image(img, file_name, path, "corners")
if self.showImages:
cv2.imshow('img', img)
cv2.waitKey(500)
# update the bar
sys.stdout.write("-")
sys.stdout.flush()
sys.stdout.write("]\n") # this ends the progress bar
print("Calibration done. Saving calibration results.")
if self.showImages:
cv2.destroyAllWindows()
# Do camera calibration given object points and image points
img_size = (img.shape[1], img.shape[0])
ret, mtx, dist, rvecs, tvecs = cv2.calibrateCamera(objpoints, imgpoints, img_size, None, None)
# Save the camera calibration result for later use (we won't worry about rvecs / tvecs)
dist_pickle = {"ret": ret, "mtx": mtx, "dist": dist, "rvecs": rvecs, "tvecs": tvecs}
calibration_file = os.path.join(self.path + "_results", "calibration_results_pickle.p")
check_dir(calibration_file)
pickle.dump(dist_pickle, open(calibration_file, "wb"))
return ret, mtx, dist, rvecs, tvecs
def predict():
"""Predict function."""
args = get_args("predict")
G_A = get_generator(args)
G_B = get_generator(args)
# Use BatchNorm2d with batchsize=1, affine=False, training=True instead of InstanceNorm2d
# Use real mean and varance rather than moving_men and moving_varance in BatchNorm2d
G_A.set_train(True)
G_B.set_train(True)
load_ckpt(args, G_A, G_B)
imgs_out = os.path.join(args.outputs_dir, "predict")
if not os.path.exists(imgs_out):
os.makedirs(imgs_out)
if not os.path.exists(os.path.join(imgs_out, "fake_A")):
os.makedirs(os.path.join(imgs_out, "fake_A"))
if not os.path.exists(os.path.join(imgs_out, "fake_B")):
os.makedirs(os.path.join(imgs_out, "fake_B"))
args.data_dir = 'testA'
ds = create_dataset(args)
reporter = Reporter(args)
reporter.start_predict("A to B")
for data in ds.create_dict_iterator(output_numpy=True):
img_A = Tensor(data["image"])
path_A = str(data["image_name"][0], encoding="utf-8")
fake_B = G_A(img_A)
save_image(fake_B, os.path.join(imgs_out, "fake_B", path_A))
reporter.info('save fake_B at %s', os.path.join(imgs_out, "fake_B",
path_A))
reporter.end_predict()
args.data_dir = 'testB'
ds = create_dataset(args)
reporter.dataset_size = args.dataset_size
reporter.start_predict("B to A")
for data in ds.create_dict_iterator(output_numpy=True):
img_B = Tensor(data["image"])
path_B = str(data["image_name"][0], encoding="utf-8")
fake_A = G_B(img_B)
save_image(fake_A, os.path.join(imgs_out, "fake_A", path_B))
reporter.info('save fake_A at %s', os.path.join(imgs_out, "fake_A",
path_B))
reporter.end_predict()
def stylize(args):
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# Loads image
content = load_image(args.content).to(device)
with torch.no_grad():
# Load Transformer net
style_model = TransformerNet()
state_dict = torch.load(args.model)
# Load Model Weights
style_model.load_state_dict(state_dict)
style_model.eval().to(device)
# Forward through Image Transformation Network
out = style_model(content).cpu()
# Save result image
save_image(out, args.out)
def saveSideImages(self, img, perpImage, file_name, data_dir, src, dst,
extension):
f, ax = plt.subplots(1, 2, figsize=(14, 5))
ax[0].imshow(img, cmap='gray')
ax[0].set_title('Image with source points drawn')
ax[1].imshow(perpImage, cmap='gray')
ax[1].set_title('Warped result with dest. points drawn')
ax[0].axis('off')
ax[1].axis('off')
x = [src[0][0], src[1][0], src[2][0], src[3][0], src[0][0]]
y = [src[0][1], src[1][1], src[2][1], src[3][1], src[0][1]]
x_ = [dst[0][0], dst[1][0], dst[2][0], dst[3][0], dst[0][0]]
y_ = [dst[0][1], dst[1][1], dst[2][1], dst[3][1], dst[0][1]]
ax[0].plot(x, y, 'b--', lw=2)
ax[1].plot(x_, y_, 'b--', lw=2)
f.tight_layout()
f.savefig("test.jpg")
save_image(mpimg.imread("test.jpg"), file_name, data_dir,
"perpSide" + extension)
os.remove("test.jpg")
def test(args):
"""Stylize a content image"""
device = torch.device("cuda" if args.cuda else "cpu")
transformer = TransformerNet().to(device)
if args.model:
transformer.load_state_dict(torch.load(args.model))
content_transform = transforms.Compose([
transforms.Resize(args.content_size),
transforms.CenterCrop(args.content_size),
transforms.ToTensor(),
transforms.Lambda(lambda x: x.mul(255))
])
content_image = utils.load_image(args.content_image)
content_image = content_transform(content_image)
content_image = content_image.unsqueeze(0).to(device)
output = transformer(content_image).cpu().detach()
utils.save_image(args.output_image, output[0] * 255)
def image_edit(args):
image_path = args.data_path
# args.save_dir = "D:/Pycharm PRoject/open_cv/data/save/images"
ch = convex_Hull()
sd = ShapeDetector()
unified = []
image = load_image(image_path)
blurred = filter_image(image)
mask = hsv_mask(blurred)
contours = find_contours(mask)
status = ch.contours_status(contours)
contour_list = ch.contour_array(status)
hull = ch.convex_hull(contour_list, contours)
unified.append(hull)
unified = np.array(unified, dtype=np.int32)
shape = sd.detect(hull)
sd.print_shape_parameters(shape)
# print("Shape", shape)
area = contour_avg_area(contours)
if area >= 15:
width = 5
elif area < 15:
width = 1
draw_contours(image, unified, width=width)
if args.save_dir is not None:
if not os.path.exists(args.save_dir):
print(
"The directory: %s doesnot exist,........ making directory: %s"
% (args.save_dir, args.save_dir))
os.makedirs(args.save_dir)
save_image(image,
args.save_dir + "/" + args.save_image,
unified,
width=width)
def undistortChessboardImages(self, data_dir):
images = glob.glob(data_dir + "/*")
for idx, file_name in enumerate(images):
img = mpimg.imread(file_name)
undistorted_img = self.undistort(data_dir, file_name, img)
save_image(undistorted_img, file_name, data_dir, "undistorted")
def test(model, data_loader, num_train_batches, epoch, test_mloss, test_rloss,
test_acc, directory):
"""
Evaluate model on validation set
Args:
model: The CapsuleNet model.
data_loader: An interator over the dataset. It combines a dataset and a sampler.
"""
print('===> Evaluate mode')
# Switch to evaluate mode
model.eval()
if args.cuda:
# When we wrap a Module in DataParallel for multi-GPUs
model = model.module
loss = 0
margin_loss = 0
recon_loss = 0
correct = 0
num_batches = len(data_loader)
global_step = epoch * num_train_batches + num_train_batches
start_time = timer()
for data, target in data_loader:
with torch.no_grad():
batch_size = data.size(0)
target_indices = target
target_one_hot = utils.one_hot_encode(target_indices,
length=args.num_classes)
assert target_one_hot.size() == torch.Size([batch_size, 10])
target = target_one_hot
if args.cuda:
data = data.to(args.device)
target = target.to(args.device)
target_indices.to(args.device)
# Output predictions
output, reconstruction = model(
data, target_indices,
False) # output from DigitCaps (out_digit_caps)
# Sum up batch loss
t_loss, m_loss, r_loss = loss_func(output, target,
args.regularization_scale,
reconstruction, data,
args.device, batch_size)
loss += t_loss.data
margin_loss += m_loss.data
recon_loss += r_loss.data
# Count number of correct predictions
# v_magnitude shape: [128, 10, 1, 1]
v_magnitude = torch.sqrt((output**2).sum(dim=2, keepdim=True))
# pred shape: [128, 1, 1, 1]
pred = v_magnitude.data.max(1, keepdim=True)[1].cpu()
correct += pred.eq(target_indices.view_as(pred)).sum()
# Get the reconstructed images of the last batch
if args.use_reconstruction_loss:
reconstruction = model.decoder(output, target_indices, False)
# Input image size and number of channel.
# By default, for MNIST, the image width and height is 28x28 and 1 channel for black/white.
image_width = args.input_width
image_height = args.input_height
image_channel = args.num_conv_in_channels
recon_img = reconstruction.view(-1, image_channel, image_width,
image_height)
assert recon_img.size() == torch.Size(
[batch_size, image_channel, image_width, image_height])
# Save the image into file system
utils.save_image(
recon_img, directory /
('recons_image_test_{}_{}.png'.format(epoch, global_step)))
utils.save_image(
data, directory /
('original_image_test_{}_{}.png'.format(epoch, global_step)))
end_time = timer()
# Log test losses
loss /= num_batches
margin_loss /= num_batches
recon_loss /= num_batches
# Log test accuracies
num_test_data = len(data_loader.dataset)
accuracy = correct / num_test_data
accuracy_percentage = float(correct) * 100.0 / float(num_test_data)
test_mloss.write('%.6f \n' % margin_loss)
test_rloss.write('%.6f \n' % recon_loss)
test_acc.write('%.4f \n' % accuracy_percentage)
# Print test losses and accuracy
print('Test: [Loss: {:.6f},' \
'\tMargin loss: {:.6f},' \
'\tReconstruction loss: {:.6f}]'.format(
loss,
margin_loss,
recon_loss if args.use_reconstruction_loss else 0))
print('Test Accuracy: {}/{} ({:.2f}%)\n'.format(correct, num_test_data,
accuracy_percentage))
global avg_testing_time_per_epoch
avg_testing_time_per_epoch = (avg_testing_time_per_epoch *
(epoch - 1) + end_time - start_time) / epoch
global best_acc
global best_acc_epoch
if accuracy_percentage > best_acc:
best_acc = accuracy_percentage
best_acc_epoch = epoch
test_loader = data_loader
utils.dump(utils.make_full_checkpoint_obj(locals(), globals()),
directory / 'trained_model/FP32_model')