class Classify(BaseInference):
def __init__(self, cfg_path, gpu_id, config_path=None):
super().__init__(config_path)
self.set_task_name(TaskName.Classify_Task)
self.task_config = self.config_factory.get_config(
self.task_name, self.config_path)
self.torchModelProcess = TorchModelProcess()
self.model = self.torchModelProcess.initModel(cfg_path, gpu_id)
self.device = self.torchModelProcess.getDevice()
def load_weights(self, weights_path):
self.torchModelProcess.loadLatestModelWeight(weights_path, self.model)
self.model = self.torchModelProcess.modelTestInit(self.model)
self.model.eval()
def process(self, input_path):
dataloader = self.get_image_data_lodaer(input_path,
self.task_config.image_size)
def infer(self, input_data, threshold=0.0):
with torch.no_grad():
output_list = self.model(input_data.to(self.device))
output = self.compute_output(output_list)
return output
def postprocess(self, result):
class_indices = torch.argmax(result, dim=1)
return class_indices
def compute_output(self, output_list):
output = None
if len(output_list) == 1:
output = self.model.lossList[0](output_list[0])
return output
class PointCloudClassify(BaseInference):
def __init__(self, cfg_path, gpu_id, config_path=None):
super().__init__(config_path)
self.set_task_name(TaskName.PC_Classify_Task)
self.task_config = self.config_factory.get_config(
self.task_name, self.config_path)
self.torchModelProcess = TorchModelProcess()
self.model = self.torchModelProcess.initModel(cfg_path, gpu_id)
self.device = self.torchModelProcess.getDevice()
def load_weights(self, weights_path):
self.torchModelProcess.loadLatestModelWeight(weights_path, self.model)
self.model = self.torchModelProcess.modelTestInit(self.model)
self.model.eval()
def process(self, input_path):
pass
def infer(self, input_data, threshold=0.0):
with torch.no_grad():
output_list = self.model(input_data.to(self.device))
output = self.compute_output(output_list)
return output
def postprocess(self, result):
pass
def compute_output(self, output_list):
output = None
loss_count = len(self.model.lossList)
if loss_count == 1:
output = self.model.lossList[0](output_list)
return output
class SuperResolution():
def __init__(self):
self.torchModelProcess = TorchModelProcess()
self.device = self.torchModelProcess.getDevice()
self.model = MSRResNet(
super_resolution_config.in_nc,
upscale_factor=super_resolution_config.upscale_factor).to(
self.device)
def load_weights(self, weights_path):
self.torchModelProcess.loadLatestModelWeight(weights_path, self.model)
self.torchModelProcess.modelTestInit(self.model)
def super_resolution(self, input_path):
if os.path.isdir(input_path):
dataloader = ImagesLoader(input_path)
else:
dataloader = VideoLoader(input_path)
prev_time = time.time()
for i, (oriImg, imgs) in enumerate(dataloader):
img_pil = Image.fromarray(cv2.cvtColor(oriImg, cv2.COLOR_BGR2RGB))
img = img_pil.convert('YCbCr')
y, cb, cr = img.split()
img_to_tensor = ToTensor()
input = img_to_tensor(y).view(1, -1, y.size[1], y.size[0])
# Get detections
with torch.no_grad():
output = self.model(input.to(self.device))[0]
print('Batch %d... Done. (%.3fs)' % (i, time.time() - prev_time))
prev_time = time.time()
out_img_y = output.cpu().detach().numpy()
out_img_y *= 255.0
out_img_y = out_img_y.clip(0, 255)
out_img_y = Image.fromarray(np.uint8(out_img_y[0]), mode='L')
out_img_cb = cb.resize(out_img_y.size, Image.BICUBIC)
out_img_cr = cr.resize(out_img_y.size, Image.BICUBIC)
out_img = Image.merge(
'YCbCr', [out_img_y, out_img_cb, out_img_cr]).convert('RGB')
show_img = cv2.cvtColor(np.asarray(out_img), cv2.COLOR_RGB2BGR)
cv2.namedWindow("image", 0)
cv2.resizeWindow("image", int(show_img.shape[1] * 0.5),
int(show_img.shape[0] * 0.5))
cv2.imshow('image', show_img)
if cv2.waitKey() & 0xFF == 27:
break
class Segmentation(BaseInference):
def __init__(self, cfg_path, gpu_id, config_path=None):
super().__init__(config_path)
self.set_task_name(TaskName.Segment_Task)
self.task_config = self.config_factory.get_config(
self.task_name, self.config_path)
self.torchModelProcess = TorchModelProcess()
self.model = self.torchModelProcess.initModel(cfg_path, gpu_id)
self.device = self.torchModelProcess.getDevice()
self.result_process = SegmentResultProcess()
self.result_show = SegmentionShow()
self.threshold = 0.5
self.src_size = (0, 0)
def load_weights(self, weights_path):
self.torchModelProcess.loadLatestModelWeight(weights_path, self.model)
self.model = self.torchModelProcess.modelTestInit(self.model)
self.model.eval()
def process(self, input_path):
dataloader = self.get_image_data_lodaer(input_path,
self.task_config.image_size)
for index, (src_image, image) in enumerate(dataloader):
self.timer.tic()
self.set_src_size(src_image)
prediction, _ = self.infer(image, self.threshold)
result = self.postprocess(prediction)
print('Batch %d... Done. (%.3fs)' % (index, self.timer.toc()))
if not self.result_show.show(src_image, result,
self.task_config.label_is_gray,
self.task_config.class_name):
break
def infer(self, input_data, threshold=0.0):
with torch.no_grad():
output_list = self.model(input_data.to(self.device))
output = self.compute_output(output_list[:])
prediction = self.result_process.get_detection_result(
output, threshold)
return prediction, output_list
def postprocess(self, result):
result = self.result_process.resize_segmention_result(
self.src_size, self.task_config.image_size, result)
return result
def compute_output(self, output_list):
count = len(output_list)
preds = []
for i in range(0, count):
temp = self.model.lossList[i](output_list[i])
preds.append(temp)
prediction = torch.cat(preds, 1)
prediction = np.squeeze(prediction.data.cpu().numpy())
return prediction
def set_src_size(self, src_data):
shape = src_data.shape[:2] # shape = [height, width]
self.src_size = (shape[1], shape[0])
class SegmentionTrain(BaseTrain):
def __init__(self, cfg_path, gpu_id, config_path=None):
super().__init__(config_path)
self.set_task_name(TaskName.Segment_Task)
self.train_task_config = self.config_factory.get_config(
self.task_name, self.config_path)
self.train_logger = TrainLogger(self.train_task_config.log_name,
self.train_task_config.root_save_dir)
self.torchModelProcess = TorchModelProcess()
self.freeze_normalization = TorchFreezeNormalization()
self.torchOptimizer = TorchOptimizer(
self.train_task_config.optimizer_config)
self.model = self.torchModelProcess.initModel(cfg_path, gpu_id)
self.device = self.torchModelProcess.getDevice()
self.output_process = SegmentResultProcess()
self.segment_test = SegmentionTest(cfg_path, gpu_id, config_path)
self.total_images = 0
self.optimizer = None
self.start_epoch = 0
self.bestmIoU = 0
def load_pretrain_model(self, weights_path):
self.torchModelProcess.loadPretainModel(weights_path, self.model)
def load_latest_param(self, latest_weights_path):
checkpoint = None
if latest_weights_path and os.path.exists(latest_weights_path):
checkpoint = self.torchModelProcess.loadLatestModelWeight(
latest_weights_path, self.model)
self.torchModelProcess.modelTrainInit(self.model)
else:
self.torchModelProcess.modelTrainInit(self.model)
self.start_epoch, self.bestmIoU = self.torchModelProcess.getLatestModelValue(
checkpoint)
self.torchOptimizer.freeze_optimizer_layer(
self.start_epoch, self.train_task_config.base_lr, self.model,
self.train_task_config.freeze_layer_name,
self.train_task_config.freeze_layer_type)
self.torchOptimizer.print_freeze_layer(self.model)
self.optimizer = self.torchOptimizer.getLatestModelOptimizer(
checkpoint)
def train(self, train_path, val_path):
dataloader = get_segment_train_dataloader(
train_path,
self.train_task_config.image_size,
self.train_task_config.train_batch_size,
is_augment=self.train_task_config.train_data_augment)
self.total_images = len(dataloader)
self.load_latest_param(self.train_task_config.latest_weights_file)
lr_factory = LrSchedulerFactory(self.train_task_config.base_lr,
self.train_task_config.max_epochs,
self.total_images)
lr_scheduler = lr_factory.get_lr_scheduler(
self.train_task_config.lr_scheduler_config)
self.train_task_config.save_config()
self.timer.tic()
self.set_model_train()
for epoch in range(self.start_epoch,
self.train_task_config.max_epochs):
# self.optimizer = torchOptimizer.adjust_optimizer(epoch, lr)
self.optimizer.zero_grad()
for idx, (images, segments) in enumerate(dataloader):
current_idx = epoch * self.total_images + idx
lr = lr_scheduler.get_lr(epoch, current_idx)
lr_scheduler.adjust_learning_rate(self.optimizer, lr)
loss = self.compute_backward(images, segments, idx)
self.update_logger(idx, self.total_images, epoch, loss)
save_model_path = self.save_train_model(epoch)
self.test(val_path, epoch, save_model_path)
def compute_backward(self, input_datas, targets, setp_index):
# Compute loss, compute gradient, update parameters
output_list = self.model(input_datas.to(self.device))
loss = self.compute_loss(output_list, targets)
loss.backward()
# accumulate gradient for x batches before optimizing
if ((setp_index + 1) % self.train_task_config.accumulated_batches == 0) \
or (setp_index == self.total_images - 1):
self.optimizer.step()
self.optimizer.zero_grad()
return loss
def compute_loss(self, output_list, targets):
loss = 0
loss_count = len(self.model.lossList)
output_count = len(output_list)
targets = targets.to(self.device)
if loss_count == 1 and output_count == 1:
output, target = self.output_process.output_feature_map_resize(
output_list[0], targets)
loss = self.model.lossList[0](output, target)
elif loss_count == 1 and output_count > 1:
loss = self.model.lossList[0](output_list, targets)
elif loss_count > 1 and loss_count == output_count:
for k in range(0, loss_count):
output, target = self.output_process.output_feature_map_resize(
output_list[k], targets)
loss += self.model.lossList[k](output, target)
else:
print("compute loss error")
return loss
def update_logger(self, index, total, epoch, loss):
loss_value = loss.data.cpu().squeeze()
step = epoch * total + index
lr = self.optimizer.param_groups[0]['lr']
self.train_logger.train_log(step, loss_value,
self.train_task_config.display)
self.train_logger.lr_log(step, lr, self.train_task_config.display)
print('Epoch: {}[{}/{}]\t Loss: {}\t Rate: {} \t Time: {}\t'.format(
epoch, index, total, '%.7f' % loss_value, '%.7f' % lr,
'%.5f' % self.timer.toc(True)))
def save_train_model(self, epoch):
self.train_logger.epoch_train_log(epoch)
if self.train_task_config.is_save_epoch_model:
save_model_path = os.path.join(
self.train_task_config.snapshot_path,
"seg_model_epoch_%d.pt" % epoch)
else:
save_model_path = self.train_task_config.latest_weights_file
self.torchModelProcess.saveLatestModel(save_model_path, self.model,
self.optimizer, epoch,
self.bestmIoU)
return save_model_path
def set_model_train(self):
self.model.train()
self.freeze_normalization.freeze_normalization_layer(
self.model, self.train_task_config.freeze_bn_layer_name,
self.train_task_config.freeze_bn_type)
def test(self, val_path, epoch, save_model_path):
if val_path is not None and os.path.exists(val_path):
self.segment_test.load_weights(save_model_path)
score, class_score, average_loss = self.segment_test.test(val_path)
self.segment_test.save_test_value(epoch, score, class_score)
self.train_logger.eval_log("val epoch loss", epoch, average_loss)
print("Val epoch loss: {}".format(average_loss))
# save best model
self.bestmIoU = self.torchModelProcess.saveBestModel(
score['Mean IoU : \t'], save_model_path,
self.train_task_config.best_weights_file)
else:
print("no test!")
class Detection2dTrain(BaseTrain):
def __init__(self, cfg_path, gpu_id, config_path=None):
super().__init__(config_path)
self.set_task_name(TaskName.Detect2d_Task)
self.train_task_config = self.config_factory.get_config(
self.task_name, self.config_path)
self.train_logger = TrainLogger(self.train_task_config.log_name,
self.train_task_config.root_save_dir)
self.torchModelProcess = TorchModelProcess()
self.freeze_normalization = TorchFreezeNormalization()
self.torchOptimizer = TorchOptimizer(
self.train_task_config.optimizer_config)
self.model = self.torchModelProcess.initModel(cfg_path, gpu_id)
self.device = self.torchModelProcess.getDevice()
self.detect_test = Detection2dTest(cfg_path, gpu_id, config_path)
self.total_images = 0
self.optimizer = None
self.avg_loss = -1
self.start_epoch = 0
self.best_mAP = 0
def load_pretrain_model(self, weights_path):
self.torchModelProcess.loadPretainModel(weights_path, self.model)
def load_latest_param(self, latest_weights_path):
checkpoint = None
if latest_weights_path and os.path.exists(latest_weights_path):
checkpoint = self.torchModelProcess.loadLatestModelWeight(
latest_weights_path, self.model)
self.model = self.torchModelProcess.modelTrainInit(self.model)
else:
self.model = self.torchModelProcess.modelTrainInit(self.model)
self.start_epoch, self.best_mAP = self.torchModelProcess.getLatestModelValue(
checkpoint)
self.torchOptimizer.freeze_optimizer_layer(
self.start_epoch, self.train_task_config.base_lr, self.model,
self.train_task_config.freeze_layer_name,
self.train_task_config.freeze_layer_type)
self.optimizer = self.torchOptimizer.getLatestModelOptimizer(
checkpoint)
def train(self, train_path, val_path):
dataloader = DetectionTrainDataloader(
train_path,
self.train_task_config.class_name,
self.train_task_config.train_batch_size,
self.train_task_config.image_size,
multi_scale=self.train_task_config.train_multi_scale,
is_augment=self.train_task_config.train_data_augment,
balanced_sample=self.train_task_config.balanced_sample)
self.total_images = len(dataloader)
lr_factory = LrSchedulerFactory(self.train_task_config.base_lr,
self.train_task_config.max_epochs,
self.total_images)
lr_scheduler = lr_factory.get_lr_scheduler(
self.train_task_config.lr_scheduler_config)
self.load_latest_param(self.train_task_config.latest_weights_file)
self.train_task_config.save_config()
self.timer.tic()
self.model.train()
self.freeze_normalization.freeze_normalization_layer(
self.model, self.train_task_config.freeze_bn_layer_name,
self.train_task_config.freeze_bn_type)
for epoch in range(self.start_epoch,
self.train_task_config.max_epochs):
# self.optimizer = self.torchOptimizer.adjust_optimizer(epoch, lr)
self.optimizer.zero_grad()
for i, (images, targets) in enumerate(dataloader):
current_iter = epoch * self.total_images + i
lr = lr_scheduler.get_lr(epoch, current_iter)
lr_scheduler.adjust_learning_rate(self.optimizer, lr)
if sum([len(x)
for x in targets]) < 1: # if no targets continue
continue
loss = self.compute_backward(images, targets, i)
self.update_logger(i, self.total_images, epoch, loss)
save_model_path = self.save_train_model(epoch)
self.test(val_path, epoch, save_model_path)
def compute_backward(self, input_datas, targets, setp_index):
# Compute loss, compute gradient, update parameters
output_list = self.model(input_datas.to(self.device))
loss = self.compute_loss(output_list, targets)
loss.backward()
# accumulate gradient for x batches before optimizing
if ((setp_index + 1) % self.train_task_config.accumulated_batches == 0) \
or (setp_index == self.total_images - 1):
self.optimizer.step()
self.optimizer.zero_grad()
return loss
def compute_loss(self, output_list, targets):
loss = 0
loss_count = len(self.model.lossList)
output_count = len(output_list)
if loss_count == 1 and output_count == 1:
loss = self.model.lossList[0](output_list[0], targets)
elif loss_count == 1 and output_count > 1:
loss = self.model.lossList[0](output_list, targets)
elif loss_count > 1 and loss_count == output_count:
for k in range(0, loss_count):
loss += self.model.lossList[k](output_list[k], targets)
else:
print("compute loss error")
return loss
def update_logger(self, index, total, epoch, loss):
step = epoch * total + index
lr = self.optimizer.param_groups[0]['lr']
loss_value = loss.data.cpu().squeeze()
if self.avg_loss < 0:
self.avg_loss = (loss.cpu().detach().numpy() /
self.train_task_config.train_batch_size)
self.avg_loss = 0.9 * (loss.cpu().detach().numpy() / self.train_task_config.train_batch_size) \
+ 0.1 * self.avg_loss
self.train_logger.train_log(step, loss_value,
self.train_task_config.display)
self.train_logger.lr_log(step, lr, self.train_task_config.display)
print('Epoch: {}[{}/{}]\t Loss: {}\t Rate: {} \t Time: {}\t'.format(
epoch, index, total, '%.7f' % self.avg_loss, '%.7f' % lr,
'%.5f' % self.timer.toc(True)))
def save_train_model(self, epoch):
self.train_logger.epoch_train_log(epoch)
if self.train_task_config.is_save_epoch_model:
save_model_path = os.path.join(
self.train_task_config.snapshot_path,
"det2d_model_epoch_%d.pt" % epoch)
else:
save_model_path = self.train_task_config.latest_weights_file
self.torchModelProcess.saveLatestModel(save_model_path, self.model,
self.optimizer, epoch,
self.best_mAP)
return save_model_path
def test(self, val_path, epoch, save_model_path):
if val_path is not None and os.path.exists(val_path):
self.detect_test.load_weights(save_model_path)
mAP, aps = self.detect_test.test(val_path)
self.detect_test.save_test_value(epoch, mAP, aps)
# save best model
self.best_mAP = self.torchModelProcess.saveBestModel(
mAP, save_model_path, self.train_task_config.best_weights_file)
else:
print("no test!")
class PointCloudClassifyTrain(BaseTrain):
def __init__(self, cfg_path, gpu_id, config_path=None):
super().__init__(config_path)
self.set_task_name(TaskName.PC_Classify_Task)
self.train_task_config = self.config_factory.get_config(
self.task_name, self.config_path)
self.train_logger = TrainLogger(self.train_task_config.log_name,
self.train_task_config.root_save_dir)
self.torchModelProcess = TorchModelProcess()
self.freeze_normalization = TorchFreezeNormalization()
self.torchOptimizer = TorchOptimizer(
self.train_task_config.optimizer_config)
self.model = self.torchModelProcess.initModel(cfg_path, gpu_id)
self.device = self.torchModelProcess.getDevice()
self.classify_test = PointCloudClassifyTest(cfg_path, gpu_id,
config_path)
self.total_clouds = 0
self.start_epoch = 0
self.best_precision = 0
self.optimizer = None
def load_pretrain_model(self, weights_path):
self.torchModelProcess.loadPretainModel(weights_path, self.model)
def load_latest_param(self, latest_weights_path):
checkpoint = None
if latest_weights_path is not None and os.path.exists(
latest_weights_path):
checkpoint = self.torchModelProcess.loadLatestModelWeight(
latest_weights_path, self.model)
self.model = self.torchModelProcess.modelTrainInit(self.model)
else:
self.model = self.torchModelProcess.modelTrainInit(self.model)
self.start_epoch, self.best_precision = self.torchModelProcess.getLatestModelValue(
checkpoint)
self.torchOptimizer.freeze_optimizer_layer(
self.start_epoch, self.train_task_config.base_lr, self.model,
self.train_task_config.freeze_layer_name,
self.train_task_config.freeze_layer_type)
self.optimizer = self.torchOptimizer.getLatestModelOptimizer(
checkpoint)
def train(self, train_path, val_path):
dataloader = get_classify_train_dataloader(
train_path, self.train_task_config.number_point_features,
self.train_task_config.train_batch_size,
self.train_task_config.train_data_augment)
self.total_clouds = len(dataloader)
lr_factory = LrSchedulerFactory(self.train_task_config.base_lr,
self.train_task_config.max_epochs,
self.total_clouds)
lr_scheduler = lr_factory.get_lr_scheduler(
self.train_task_config.lr_scheduler_config)
self.load_latest_param(self.train_task_config.latest_weights_file)
self.train_task_config.save_config()
self.timer.tic()
self.model.train()
self.freeze_normalization.freeze_normalization_layer(
self.model, self.train_task_config.freeze_bn_layer_name,
self.train_task_config.freeze_bn_type)
try:
for epoch in range(self.start_epoch,
self.train_task_config.max_epochs):
# self.optimizer = torchOptimizer.adjust_optimizer(epoch, lr)
self.optimizer.zero_grad()
for idx, (clouds, targets) in enumerate(dataloader):
current_iter = epoch * self.total_clouds + idx
lr = lr_scheduler.get_lr(epoch, current_iter)
lr_scheduler.adjust_learning_rate(self.optimizer, lr)
loss = self.compute_backward(clouds, targets, idx)
self.update_logger(idx, self.total_clouds, epoch, loss)
save_model_path = self.save_train_model(epoch)
self.test(val_path, epoch, save_model_path)
except Exception as e:
raise e
finally:
self.train_logger.close()
def compute_backward(self, input_datas, targets, setp_index):
# Compute loss, compute gradient, update parameters
output_list = self.model(input_datas.to(self.device))
loss = self.compute_loss(output_list, targets)
loss.backward()
# accumulate gradient for x batches before optimizing
if ((setp_index + 1) % self.train_task_config.accumulated_batches == 0) or \
(setp_index == self.total_clouds - 1):
self.optimizer.step()
self.optimizer.zero_grad()
return loss
def compute_loss(self, output_list, targets):
loss = 0
loss_count = len(self.model.lossList)
output_count = len(output_list)
targets = targets.to(self.device)
if loss_count == 1 and output_count == 1:
loss = self.model.lossList[0](output_list[0], targets)
elif loss_count == 1 and output_count > 1:
loss = self.model.lossList[0](output_list, targets)
elif loss_count > 1 and loss_count == output_count:
for k in range(0, loss_count):
loss += self.model.lossList[k](output_list[k], targets)
else:
print("compute loss error")
return loss
def update_logger(self, index, total, epoch, loss):
step = epoch * total + index
lr = self.optimizer.param_groups[0]['lr']
loss_value = loss.data.cpu().squeeze()
self.train_logger.train_log(step, loss_value,
self.train_task_config.display)
self.train_logger.lr_log(step, lr, self.train_task_config.display)
print('Epoch: {}[{}/{}]\t Loss: {}\t Rate: {} \t Time: {}\t'.format(
epoch, index, total, '%.7f' % loss_value, '%.7f' % lr,
'%.5f' % self.timer.toc(True)))
def save_train_model(self, epoch):
with DelayedKeyboardInterrupt():
self.train_logger.epoch_train_log(epoch)
if self.train_task_config.is_save_epoch_model:
save_model_path = os.path.join(
self.train_task_config.snapshot_path,
"pc_cls_model_epoch_%d.pt" % epoch)
else:
save_model_path = self.train_task_config.latest_weights_file
self.torchModelProcess.saveLatestModel(save_model_path, self.model,
self.optimizer, epoch,
self.best_precision)
return save_model_path
def test(self, val_path, epoch, save_model_path):
if val_path is not None and os.path.exists(val_path):
self.classify_test.load_weights(save_model_path)
precision = self.classify_test.test(val_path)
self.classify_test.save_test_value(epoch)
# save best model
self.best_precision = self.torchModelProcess.saveBestModel(
precision, save_model_path,
self.train_task_config.best_weights_file)
else:
print("no test!")
class Detection2d(BaseInference):
def __init__(self, cfg_path, gpu_id, config_path=None):
super().__init__(config_path)
self.set_task_name(TaskName.Detect2d_Task)
self.task_config = self.config_factory.get_config(
self.task_name, self.config_path)
self.torchModelProcess = TorchModelProcess()
self.result_process = Detect2dResultProcess()
self.nms_process = FastNonMaxSuppression()
self.result_show = DetectionShow()
self.model = self.torchModelProcess.initModel(cfg_path, gpu_id)
self.device = self.torchModelProcess.getDevice()
self.src_size = (0, 0)
def load_weights(self, weights_path):
self.torchModelProcess.loadLatestModelWeight(weights_path, self.model)
self.model = self.torchModelProcess.modelTestInit(self.model)
self.model.eval()
def process(self, input_path):
dataloader = self.get_image_data_lodaer(input_path,
self.task_config.image_size)
for i, (src_image, img) in enumerate(dataloader):
print('%g/%g' % (i + 1, len(dataloader)), end=' ')
self.set_src_size(src_image)
self.timer.tic()
result = self.infer(img, self.task_config.confidence_th)
detection_objects = self.postprocess(result)
print('Batch %d... Done. (%.3fs)' % (i, self.timer.toc()))
if not self.result_show.show(src_image, detection_objects):
break
def save_result(self, filename, detection_objects):
for object in detection_objects:
confidence = object.classConfidence * object.objectConfidence
x1 = object.min_corner.x
y1 = object.min_corner.y
x2 = object.max_corner.x
y2 = object.max_corner.y
temp_save_path = os.path.join(self.task_config.save_result_dir,
"%s.txt" % object.name)
with open(temp_save_path, 'a') as file:
file.write("{} {} {} {} {} {}\n".format(
filename, confidence, x1, y1, x2, y2))
def infer(self, input_data, threshold=0.0):
with torch.no_grad():
output_list = self.model(input_data.to(self.device))
output = self.compute_output(output_list)
result = self.result_process.get_detection_result(
output, threshold)
return result
def postprocess(self, result):
detection_objects = self.nms_process.multi_class_nms(
result, self.task_config.nms_th)
detection_objects = self.result_process.resize_detection_objects(
self.src_size, self.task_config.image_size, detection_objects,
self.task_config.class_name)
return detection_objects
def compute_output(self, output_list):
count = len(output_list)
preds = []
for i in range(0, count):
temp = self.model.lossList[i](output_list[i])
preds.append(temp)
prediction = torch.cat(preds, 1)
prediction = prediction.squeeze(0)
return prediction
def set_src_size(self, src_data):
shape = src_data.shape[:2] # shape = [height, width]
self.src_size = (shape[1], shape[0])
class SuperResolutionTrain():
def __init__(self):
if not os.path.exists(super_resolution_config.snapshotPath):
os.makedirs(super_resolution_config.snapshotPath, exist_ok=True)
self.torchModelProcess = TorchModelProcess()
self.torchOptimizer = TorchOptimizer(super_resolution_config.optimizerConfig)
self.multiLR = MultiStageLR(super_resolution_config.base_lr, [[50, 1], [70, 0.1], [100, 0.01]])
self.device = self.torchModelProcess.getDevice()
self.model = MSRResNet(super_resolution_config.in_nc, upscale_factor=super_resolution_config.upscale_factor).to(self.device)
vision_process = TorchVisionProcess()
self.input_transform = vision_process.input_transform(super_resolution_config.crop_size,
super_resolution_config.upscale_factor)
self.target_transform = vision_process.target_transform(super_resolution_config.crop_size)
self.start_epoch = 0
self.psnr = 0
def load_param(self, latest_weights_path):
checkpoint = None
if latest_weights_path and os.path.exists(latest_weights_path):
checkpoint = self.torchModelProcess.loadLatestModelWeight(latest_weights_path, self.model)
self.torchModelProcess.modelTrainInit(self.model)
else:
self.torchModelProcess.modelTrainInit(self.model)
self.start_epoch, self.best_mAP = self.torchModelProcess.getLatestModelValue(checkpoint)
self.torchOptimizer.createOptimizer(self.start_epoch, self.model, super_resolution_config.base_lr)
self.optimizer = self.torchOptimizer.getLatestModelOptimizer(checkpoint)
def update_logger(self, step, value):
pass
def compute_loss(self, output, targets):
loss = 0
criterion = nn.MSELoss()
loss += criterion(output, targets)
return loss
def test(self, epoch):
save_model_path = os.path.join(super_resolution_config.snapshotPath, "model_epoch_%d.pt" % epoch)
self.torchModelProcess.saveLatestModel(save_model_path, self.model,
self.optimizer, epoch, self.psnr)
testing_data_loader = get_sr_dataloader(super_resolution_config.val_set,
super_resolution_config.val_set,
super_resolution_config.test_batch_size,
num_workers=8,
input_transform=self.input_transform,
target_transform=self.target_transform)
avg_psnr = 0
with torch.no_grad():
for batch in testing_data_loader:
input, target = batch[0].to(self.device), batch[1].to(self.device)
prediction = self.model(input)
mse = self.compute_loss(prediction, target)
psnr = 10 * log10(1 / mse.item())
avg_psnr += psnr
self.psnr = avg_psnr / len(testing_data_loader)
self.torchModelProcess.saveBestModel(self.psnr,
save_model_path,
super_resolution_config.best_weights_file)
print("===> Avg. PSNR: {:.4f} dB".format(self.psnr))
def train(self):
training_data_loader = get_sr_dataloader(super_resolution_config.train_set,
super_resolution_config.train_set,
super_resolution_config.train_batch_size,
num_workers=8,
shuffle=True,
input_transform=self.input_transform,
target_transform=self.target_transform)
total_images = len(training_data_loader)
self.load_param(super_resolution_config.latest_weights_file)
t0 = time.time()
for epoch in range(self.start_epoch, super_resolution_config.maxEpochs):
self.optimizer.zero_grad()
for i, batch in enumerate(training_data_loader, 1):
imgs, targets = batch[0].to(self.device), batch[1].to(self.device)
current_iter = epoch * total_images + i
lr = self.multiLR.get_lr(epoch, current_iter)
self.multiLR.adjust_learning_rate(self.optimizer, lr)
# Compute loss, compute gradient, update parameters
output = self.model(imgs.to(self.device))
loss = self.compute_loss(output, targets)
loss.backward()
# accumulate gradient for x batches before optimizing
if ((i + 1) % super_resolution_config.accumulated_batches == 0) or (i == len(training_data_loader) - 1):
self.optimizer.step()
self.optimizer.zero_grad()
print('Epoch: {}[{}/{}]\t Loss: {}\t Rate: {} \t Time: {}\t'.format(epoch, i, len(training_data_loader),
'%.3f' % loss,
'%.7f' %
self.optimizer.param_groups[0][
'lr'], time.time() - t0))
self.update_logger(current_iter, loss.data)
t0 = time.time()
self.test(epoch)