From 657cfcc0963304157e67f728d0f04a36004fdd69 Mon Sep 17 00:00:00 2001 From: zpeng11 Date: Mon, 30 Dec 2024 00:09:04 -0800 Subject: [PATCH] Get interpolation ready --- args.py | 13 ++++++++++--- main.py | 35 ++++++++++++++--------------------- 2 files changed, 24 insertions(+), 24 deletions(-) diff --git a/args.py b/args.py index df01dd8..b59be68 100644 --- a/args.py +++ b/args.py @@ -43,7 +43,7 @@ parser.add_argument('--model_vram_cache', type=bool, default=True) # parser.add_argument('--model_cache_size', type=float, default=1.0) parser.add_argument('--use_interpolation', type=bool, default=False) -parser.add_argument('--interpolation_scale', type=int, default=2) +parser.add_argument('--interpolation_scale', type=int, default=3) parser.add_argument('--interpolation_half', type=bool, default=True) parser.add_argument('--use_cacher', type=bool, default=True) @@ -68,5 +68,12 @@ else: if args.output_webcam is None and args.output_dir is None: args.debug = True -args.device_id = 1 -args.use_tensorrt = False \ No newline at end of file + +args.device_id = 2 +args.use_tensorrt = False +args.use_interpolation = True +args.interpolation_scale = 4 +args.use_cacher = False +args.model_half = False + +args.model_output_size = 1024 if args.use_sr else 512 \ No newline at end of file diff --git a/main.py b/main.py index cf92e80..6b0eb94 100644 --- a/main.py +++ b/main.py @@ -325,7 +325,7 @@ class ModelClientProcess(Process): self.input_image = input_image self.output_queue = Queue() self.input_queue = Queue() - self.shms = [shared_memory.SharedMemory(create=True, size=512 * 512 * 4) for _ in range(1)] + self.shms = [shared_memory.SharedMemory(create=True, size=args.model_output_size * args.model_output_size * 4) for _ in range(args.interpolation_scale if args.use_interpolation else 1)] self.model_fps_number = Value('f', 0.0) self.gpu_fps_number = Value('f', 0.0) self.cache_hit_ratio = Value('f', 0.0) @@ -356,7 +356,7 @@ class ModelClientProcess(Process): self.model.setImage(self.input_image) input_pose = np.zeros((1,45), dtype=np.float32) - self.shared_nps = [np.ndarray((512, 512, 4), dtype=np.uint8, buffer=self.shms[i].buf) for i in range(len(self.shms))] + self.shared_nps = [np.ndarray((args.model_output_size, args.model_output_size, 4), dtype=np.uint8, buffer=self.shms[i].buf) for i in range(len(self.shms))] model_fps = FPS() while True: @@ -615,7 +615,8 @@ def main(): model_process = ModelClientProcess(input_image) model_process.daemon = True model_process.start() - model_output_nps = [np.ndarray((512, 512, 4), dtype=np.uint8, buffer=model_process.shms[i].buf) for i in range(len(model_process.shms))] + model_output_nps = [np.ndarray((args.model_output_size, args.model_output_size, 4), dtype=np.uint8, buffer=model_process.shms[i].buf) for i in range(len(model_process.shms))] + model_result_read_ptr = 0 print("Ready. Close this console to exit.") @@ -831,23 +832,19 @@ def main(): model_process.input_queue.put_nowait(model_input_arr) - has_model_output = 0 + try: - new_model_output = model_output while not model_process.output_queue.empty(): - has_model_output += 1 - new_model_output = model_process.output_queue.get_nowait() - model_output = model_output_nps[0].copy() + model_process.output_queue.get_nowait() + model_result_read_ptr = 0 + model_output = model_output_nps[model_result_read_ptr].copy() + if model_result_read_ptr + 1 < len(model_output_nps): + model_result_read_ptr += 1 except queue.Empty: pass if model_output is None: time.sleep(1) continue - # print(has_model_output) - # should_output=should_output or has_model_output - # if not should_output: - # continue - postprocessed_image = model_output if args.perf == 'main': @@ -870,18 +867,14 @@ def main(): if args.bongo: rotate_angle -= 5 - new_size = postprocessed_image.shape[0] - IMG_WIDTH = new_size - args.output_w = new_size - args.output_h = new_size - rm = cv2.getRotationMatrix2D((IMG_WIDTH / 2, IMG_WIDTH / 2), rotate_angle, k_scale) - rm[0, 2] += dx + args.output_w / 2 - IMG_WIDTH / 2 - rm[1, 2] += dy + args.output_h / 2 - IMG_WIDTH / 2 + rm = cv2.getRotationMatrix2D((args.model_output_size / 2, args.model_output_size / 2), rotate_angle, k_scale) + rm[0, 2] += dx + args.output_w / 2 - args.model_output_size / 2 + rm[1, 2] += dy + args.output_h / 2 - args.model_output_size / 2 postprocessed_image = cv2.warpAffine( postprocessed_image, rm, - (args.output_w, args.output_h)) + (args.model_output_size, args.model_output_size)) if args.perf == 'main': print("extendmovement", (time.perf_counter() - tic) * 1000)