def _G_pred(self, input, now_frame, label, prev_class):
input_part = input[:, :, :, now_frame:now_frame +
self.test_frame, :, :].to(self.device1)
label_og_part = label[:, now_frame:now_frame + self.test_frame].to(
self.device1)
# tensor to numpy and label condition
label_og_part = tensor2np(label_og_part)
input_part = tensor2np(input_part)
input_part, label_part = vid_init_test(label_og_part, prev_class,
input_part)
# [V, D, H, W, C] -> [N, V, C, D, H, W]
input_part = np.expand_dims(input_part, axis=0)
label_part = np.expand_dims(label_part, axis=0)
label_og_part = np.expand_dims(label_og_part, axis=0)
# numpy to tensor
input_part = torch.from_numpy(input_part).to(self.device1)
label_part = torch.from_numpy(label_part).to(self.device1)
label_og_part = torch.from_numpy(label_og_part).to(self.device1)
features = self.extractor(input_part[:, 0, :, :, :, :])
for i in range(1, args.num_of_vid):
features = torch.cat(
(features, self.extractor(input_part[:, 0, :, :, :, :])),
dim=1)
predict = self.G(features)
return predict, label_part, label_og_part
def valid():
model.eval()
avg_psnr, avg_ssim = 0, 0
for i, batch in enumerate(testing_data_loader):
lr_tensor, hr_tensor = batch[0], batch[1]
if args.cuda:
lr_tensor = lr_tensor.to(device)
hr_tensor = hr_tensor.to(device)
with torch.no_grad():
pre = model(lr_tensor)
sr_img = utils.tensor2np(pre.detach()[0])
gt_img = utils.tensor2np(hr_tensor.detach()[0])
crop_size = args.scale
cropped_sr_img = utils.shave(sr_img, crop_size)
cropped_gt_img = utils.shave(gt_img, crop_size)
if args.isY is True:
im_label = utils.quantize(sc.rgb2ycbcr(cropped_gt_img)[:, :, 0])
im_pre = utils.quantize(sc.rgb2ycbcr(cropped_sr_img)[:, :, 0])
else:
im_label = cropped_gt_img
im_pre = cropped_sr_img
psnr = utils.compute_psnr(im_pre, im_label)
ssim = utils.compute_ssim(im_pre, im_label)
avg_psnr += psnr
avg_ssim += ssim
print(
f" Valid {i}/{len(testing_data_loader)} with PSNR = {psnr} and SSIM = {ssim}"
)
print("===> Valid. psnr: {:.4f}, ssim: {:.4f}".format(
avg_psnr / len(testing_data_loader),
avg_ssim / len(testing_data_loader)))
def _pred_label(self, og_label, changed_label, predict, k, predict_label,
i):
pred = predict[:, i, :]
_, pred = torch.max(pred.data, 1)
frame_num = k * self.test_frame + i
pred = tensor2np(pred)
ch_la = tensor2np(changed_label)
og_la = tensor2np(og_label)
# return
pred = vid_init_return(og_la, ch_la, pred)
pred = [frame_num, pred]
predict_label.append(pred)
def train(model, device, trainloader, optimizer, loss_function):
model.train()
running_loss = 0
mask_list, iou = [], []
for i, (input, mask) in enumerate(trainloader):
# load data into cuda
input, mask = input.to(device,
dtype=torch.float), mask.to(device,
dtype=torch.float)
# forward
predict = model(input)
loss = loss_function(predict, mask)
# metric
iou.append(get_iou_score(predict, mask).mean())
running_loss += (loss.item())
# zero the gradient + backprpagation + step
optimizer.zero_grad()
loss.backward()
optimizer.step()
# log the first image of the batch
if ((i + 1) % 20) == 0 and not args.tuning:
rand = np.random.randint(0, 6000)
img, pred, mak = tensor2np(input[1]), tensor2np(
predict[1]), tensor2np(mask[1])
savenp2Img(SAVE_PATH + f'image_{rand}.jpg', img)
savenp2Img(SAVE_PATH + f'prediction_{rand}.jpg', pred)
savenp2Img(SAVE_PATH + f'mask_{rand}.jpg', mak)
mask_list.extend([
wandb.Image(SAVE_PATH + f'image_{rand}.jpg'),
wandb.Image(SAVE_PATH + f'mask_{rand}.jpg'),
wandb.Image(SAVE_PATH + f'prediction_{rand}.jpg'),
])
mean_iou = np.mean(iou)
total_loss = running_loss / len(trainloader)
if not args.tuning:
wandb.log({
'Train loss': total_loss,
'Train IoU': mean_iou,
'Train prediction': mask_list
})
else:
wandb.log({'Train loss': total_loss, 'Train IoU': mean_iou})
return total_loss, mean_iou
def test(model, device, testloader, loss_function, best_iou):
model.eval()
running_loss = 0
mask_list, iou = [], []
with torch.no_grad():
for i, (input, mask) in enumerate(testloader):
input, mask = input.to(device, dtype=torch.float), mask.to(
device, dtype=torch.float)
predict = model(input)
loss = loss_function(predict, mask)
running_loss += loss.item()
iou.append(get_iou_score(predict, mask).mean())
# log the first image of the batch
if ((i + 1) % 10) == 0 and not args.tuning:
rand = np.random.randint(0, 6000)
img, pred, mak = tensor2np(input[1]), tensor2np(
predict[1]), tensor2np(mask[1])
savenp2Img(SAVE_PATH + f'image_{rand}.jpg', img)
savenp2Img(SAVE_PATH + f'prediction_{rand}.jpg', pred)
savenp2Img(SAVE_PATH + f'mask_{rand}.jpg', mak)
mask_list.extend([
wandb.Image(SAVE_PATH + f'image_{rand}.jpg'),
wandb.Image(SAVE_PATH + f'mask_{rand}.jpg'),
wandb.Image(SAVE_PATH + f'prediction_{rand}.jpg'),
])
test_loss = running_loss / len(testloader)
mean_iou = np.mean(iou)
if not args.tuning:
wandb.log({
'Valid loss': test_loss,
'Valid IoU': mean_iou,
'Prediction': mask_list
})
else:
wandb.log({'Valid loss': test_loss, 'Valid IoU': mean_iou})
if mean_iou > best_iou:
# export to onnx + pt
try:
torch.onnx.export(model, input, SAVE_PATH + RUN_NAME + '.onnx')
torch.save(model.state_dict(), SAVE_PATH + RUN_NAME + '.pth')
except:
print('Can export weights')
return test_loss, mean_iou
def sr_process(model):
# im_l = cv2.imread(imname, cv2.IMREAD_COLOR)[:, :, [2, 1, 0]] # BGR to RGB
while (True):
image = frame_queue.get()
start = time.time() * 1000
im_l = image[:, :, [2, 1, 0]]
im_input = im_l / 255.0
im_input = np.transpose(im_input, (2, 0, 1))
im_input = im_input[np.newaxis, ...]
im_input = torch.from_numpy(im_input).float()
if cuda:
im_input = im_input.to(device)
print(f"time cost 01 = {time.time()*1000-start} ms")
# with torch.no_grad():
out = model(im_input)
torch.cuda.synchronize()
print(f"time cost 02 = {time.time()*1000-start} ms")
out_img = utils.tensor2np(out.detach()[0])
print(f"time cost 03 = {time.time()*1000-start} ms")
def test(self):
with torch.no_grad():
# predict label for saving csv file
for iter, batch in enumerate(self.test_data, 1):
# for save predicted label
predict_label = []
input, label = batch[0].to(self.device0), batch[1].to(
self.device0)
vid_total_frame = label.shape[1]
total_iter = int(vid_total_frame / self.test_frame)
total_frame = int(label.shape[1] / self.test_frame)
last_iter_start = int(vid_total_frame - self.test_frame)
print('video total frames : ', vid_total_frame)
# start video condition
prev_class = self._start_vid(iter)
for k in range(total_frame):
if k != 0:
prev_class = np.array(predict_label)[-1, 1]
now_frame = k * self.test_frame
predict, label_part, label_og_part = \
self._G_pred(input, now_frame, label, prev_class)
for i in range(self.test_frame):
self._pred_label(label_og_part, label_part, predict, k,
predict_label, i)
# last iteration
prev_class = np.array(predict_label)[-1, 1]
now_frame = last_iter_start
predict, label_part, label_og_part = \
self._G_pred(input, now_frame, label, prev_class)
# append additional predict label
for i in range(self.test_frame):
frame_num = now_frame + i
if frame_num < total_iter * self.test_frame:
pass
else:
self._pred_label(label_og_part, label_part, predict, k,
predict_label, i)
# calculate accuracy
predict_label = np.array(predict_label)
label = tensor2np(label)
print("accuracy : ", self._cal_accuracy(label, predict_label))
# sequnetial processing
sequential_predict_label = self._sequential_process(
np.array(predict_label), 2)
sequential_predict_label = self._sequential_process(
np.array(sequential_predict_label), 1)
# save results video
self.save_pred_label(sequential_predict_label, iter)
# clear gpu
torch.cuda.empty_cache()
def process(image):
im_l = image[:, :, [2, 1, 0]] # BGR to RGB
# im_l = sio.imread(opt.test_lr_folder + '/' + imname.split('/')[-1]) # RGB
im_input = im_l / 255.0
im_input = np.transpose(im_input, (2, 0, 1))
im_input = im_input[np.newaxis, ...]
im_input = torch.from_numpy(im_input).float()
if cuda:
# print("use cuda!!!")
im_input = im_input.to(device)
with torch.no_grad():
start = time.time()*1000
out = model(im_input)
torch.cuda.synchronize()
# time_cost = time.time()*1000 - start
# print(f"time cost = {time_cost} ms")
out_img = utils.tensor2np(out.detach()[0])
print(f"time cost = {time.time()*1000 - start} ms")
return out_img[:, :, [2, 1, 0]]
def process(imname, model):
im_l = cv2.imread(imname, cv2.IMREAD_COLOR)[:, :, [2, 1, 0]] # BGR to RGB
# im_l = sio.imread(opt.test_lr_folder + '/' + imname.split('/')[-1]) # RGB
im_input = im_l / 255.0
im_input = np.transpose(im_input, (2, 0, 1))
im_input = im_input[np.newaxis, ...]
im_input = torch.from_numpy(im_input).float()
if cuda:
print("use cuda!!!")
im_input = im_input.to(device)
with torch.no_grad():
start = time.time() * 1000
out = model(im_input)
torch.cuda.synchronize()
time_cost = time.time() * 1000 - start
print(f"time cost = {time_cost} ms")
out_img = utils.tensor2np(out.detach()[0])
crop_size = upscale_factor
cropped_sr_img = utils.shave(out_img, crop_size)
im_input = np.transpose(im_input, (2, 0, 1))
im_input = im_input[np.newaxis, ...]
im_input = torch.from_numpy(im_input).float()
if cuda:
model = model.to(device)
im_input = im_input.to(device)
with torch.no_grad():
start.record()
out = model(im_input)
end.record()
torch.cuda.synchronize()
time_list[i] = start.elapsed_time(end) # milliseconds
out_img = utils.tensor2np(out.detach()[0])
crop_size = opt.upscale_factor
cropped_sr_img = utils.shave(out_img, crop_size)
cropped_gt_img = utils.shave(im_gt, crop_size)
if opt.is_y is True:
im_label = utils.quantize(sc.rgb2ycbcr(cropped_gt_img)[:, :, 0])
im_pre = utils.quantize(sc.rgb2ycbcr(cropped_sr_img)[:, :, 0])
else:
im_label = cropped_gt_img
im_pre = cropped_sr_img
psnr_list[i] = utils.compute_psnr(im_pre, im_label)
ssim_list[i] = utils.compute_ssim(im_pre, im_label)
output_folder = os.path.join(
opt.output_folder,
imname.split('/')[-1].split('.')[0] + 'x' + str(opt.upscale_factor) +
model.eval()
avg_psnr, avg_ssim = 0, 0
start = time.time()
for _,batch in enumerate(testing_data_loader):
lr_tensor, hr_tensor = batch[0], batch[1]
if 1:
lr_tensor = lr_tensor.to(device)
hr_tensor = hr_tensor.to(device)
with torch.no_grad():
pre = model(lr_tensor)
sr_img = utils.tensor2np(pre.detach()[0])
gt_img = utils.tensor2np(hr_tensor.detach()[0])
crop_size = 2
cropped_sr_img = utils.shave(sr_img, crop_size)
cropped_gt_img = utils.shave(gt_img, crop_size)
im_label = utils.quantize(sc.rgb2ycbcr(cropped_gt_img)[:, :, 0])
im_pre = utils.quantize(sc.rgb2ycbcr(cropped_sr_img)[:, :, 0])
img_save=sr_img.transpose(2,0,1)
cv2.imwrite(os.path.join("./results/",str(_)+".png"),img_save)
avg_psnr += utils.compute_psnr(im_pre, im_label)
avg_ssim += utils.compute_ssim(im_pre, im_label)
def test(self):
with torch.no_grad():
# predict label for saving csv file
for iter, batch in enumerate(self.test_data, 1):
# for save predicted label
predict_label = []
input, label = batch[0].to(self.device0), batch[1].to(
self.device0)
vid_total_frame = label.shape[1]
print('video total frames : ', vid_total_frame)
total_iter = int(vid_total_frame / self.test_frame)
total_frame = int(label.shape[1] / self.test_frame)
last_iter_start = int(vid_total_frame - self.test_frame)
for k in range(total_frame):
now_frame = k * self.test_frame
input_part = input[:, :, :, now_frame:now_frame +
self.test_frame, :, :].to(self.device1)
label_og_part = label[:, now_frame:now_frame +
self.test_frame].to(self.device1)
if k == 0 and iter == 1:
prev_class = 0
elif k == 0 and iter == 2:
prev_class = 6
elif k == 0 and iter == 3:
prev_class = 1
elif k == 0 and iter == 4:
prev_class = 2
else:
prev_class = np.array(predict_label)[-1, 1]
# tensor에서 numpy로 변환
label_og_part = tensor2np(label_og_part)
input_part = tensor2np(input_part)
input_part, label_part = vid_init_test(
label_og_part, prev_class, input_part)
# [V, D, H, W, C] -> [N, V, C, D, H, W]
input_part = np.expand_dims(input_part, axis=0)
label_part = np.expand_dims(label_part, axis=0)
label_og_part = np.expand_dims(label_og_part, axis=0)
# numpy에서 tesnor 변환
input_part = torch.from_numpy(input_part).to(self.device1)
label_part = torch.from_numpy(label_part).to(self.device1)
label_og_part = torch.from_numpy(label_og_part).to(
self.device1)
features = self.extractor(input_part[:, 0, :, :, :, :])
for i in range(1, args.num_of_vid):
features = torch.cat(
(features,
self.extractor(input_part[:, 0, :, :, :, :])),
dim=1)
predict = self.G(features)
for i in range(self.test_frame):
self._pred_label(label_og_part, label_part, predict, k,
predict_label, i)
# last iteration
now_frame = last_iter_start
input_part = input[:, :, :, now_frame:now_frame +
self.test_frame, :, :].to(self.device1)
label_og_part = label[:, now_frame:now_frame +
self.test_frame].to(self.device1)
# vid init을 해줘서 label_og와 label을 구분
if k == 0:
prev_class = 0
else:
prev_class = np.array(predict_label)[-1, 1]
# tensor에서 numpy로 변환
label_og_part = tensor2np(label_og_part)
input_part = tensor2np(input_part)
input_part, label_part = vid_init_test(label_og_part,
prev_class, input_part)
# [V, D, H, W, C] -> [N, V, C, D, H, W]
input_part = np.expand_dims(input_part, axis=0)
label_part = np.expand_dims(label_part, axis=0)
label_og_part = np.expand_dims(label_og_part, axis=0)
# numpy에서 tesnor 변환
input_part = torch.from_numpy(input_part).to(self.device1)
label_part = torch.from_numpy(label_part).to(self.device1)
label_og_part = torch.from_numpy(label_og_part).to(
self.device1)
features = self.extractor(input_part[:, 0, :, :, :, :])
for i in range(1, args.num_of_vid):
features = torch.cat(
(features, self.extractor(input_part[:,
0, :, :, :, :])),
dim=1)
predict = self.G(features)
# add predict label
for i in range(self.test_frame):
frame_num = now_frame + i
if frame_num < total_iter * self.test_frame:
pass
else:
self._pred_label(label_og_part, label_part, predict, k,
predict_label, i)
# sequnetial processing
sequential_predict_label = self.sequential_process(
np.array(predict_label), 2)
sequential_predict_label = self.sequential_process(
np.array(sequential_predict_label), 1)
# save results video
self.save_pred_label(sequential_predict_label, iter)
# clear gpu
torch.cuda.empty_cache()