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 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 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 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])