writer.add_scalars('loss/per_epoch', {'train_loss': train_epoch_loss / n},
epoch)
train_loss = 0
train_epoch_loss = 0
data_loader_iter = iter(data_loader)
for img, mask in data_loader_iter:
solver.set_input(img, mask)
train_loss = solver.optimize()
train_epoch_loss += train_loss
#print(train_loss)
#print(iter_count)
if (iter_count % 1000) == 0:
t_img = solver.test_one_img_to_real(img)
#print(t_img.shape)
#print(1) * 2 - 1
writer.add_image('raw_image/raw_image', t_img[0] * 255, iter_count)
t_mask = solver.test_one_img_to_real(mask)
writer.add_image('raw_mask/raw_mask', t_mask[0] * 255, iter_count)
#iter_img = solver.test_one_img(img)
iter_img = solver.test_one_mask_to_real(img) #[0,]
#print(iter_img.shape)
iter_img = iter_img[0, ]
#print(iter_img.shape)
iter_img = iter_img.reshape(1, 512, 512)
#print(iter_img.shape)
writer.add_image('gen_mask/gen_mask', iter_img, iter_count)
#print(train_loss)
h, w, c = src.shape
img = np.array(src)
dst = src.copy()
dist = []
temp = np.array(src)
temp = np.array(temp, np.float32).transpose(2,0,1)/255.0
dist.append(temp)
dist = np.array(dist)
img = torch.Tensor(dist)
solver.set_input(img)#, mask)
_ = solver.forward()#optimize()
t_img = solver.test_one_img_to_real(img)
t_img= t_img[0]
t_img = t_img.transpose(1, 2, 0)
t_img_out = t_img
t_img_out = t_img_out.transpose(2,0,1)
t_img_out = np.uint8(t_img_out.transpose(2,0,1))
#print(t_img_out.shape)
iter_img = solver.test_one_mask_to_real(img)#[0,]
iter_img = np.uint8(iter_img)
#Image.fromarray(t_img_out).save(RESULT_DIR + 'raw_' + str(iter_count) + '_' + str(j) + '.jpg')
t_out = Image.fromarray(iter_img, mode='L')
#t_out.save(RESULT_DIR + 'result_' + str(iter_count) + '_' + str(j) + '.jpg')
