diff --git a/EndoCalibrator.py b/EndoCalibrator.py index fa8baee..26eb09d 100644 --- a/EndoCalibrator.py +++ b/EndoCalibrator.py @@ -53,6 +53,7 @@ zoom_y = resize_size[1] / gray.shape[0] intrinsics_scaled = self.scale_intrinsics(mtx, zoom_x, zoom_y) os.makedirs(self.args.npy_outdir, exist_ok=True) + np.save(os.path.join(self.args.npy_outdir, 'intrinsics_row.npy'), mtx) np.save(os.path.join(self.args.npy_outdir, 'intrinsics_scaled.npy'), intrinsics_scaled) np.save(os.path.join(self.args.npy_outdir, 'dist_coeffs.npy'), dist) np.save(os.path.join(self.args.npy_outdir, 'caribrate_img_size.npy'), np.array([gray.shape[1], gray.shape[0]])) diff --git a/Movie2imgConverter.py b/Movie2imgConverter.py index f0ef2ab..309b7d7 100644 --- a/Movie2imgConverter.py +++ b/Movie2imgConverter.py @@ -6,6 +6,10 @@ class Movie2imgConverter: def __init__(self, args): self.args = args + self.calibrate_img_size = np.load(os.path.join(self.args.param_dir, 'caribrate_img_size.npy')) + self.dist_coeffs = np.load(os.path.join(self.args.param_dir, 'dist_coeffs.npy')) + self.intrinsics_row = np.load(os.path.join(self.args.param_dir, 'intrinsics_row.npy')) + self.resized_size = (self.args.resized_width, self.args.resized_height) def convert_endo_movie2img(self): src_dir = self.args.movie_dir @@ -20,7 +24,56 @@ raise sequence_num = 0 + file_name_num = 0 for movie_name in movie_name_list: + cap = cv2.VideoCapture(movie_name) + if not cap.isOpened(): + print("ビデオキャプチャエラー") + raise + print("start processing {}".format(os.path.basename(movie_name))) - sum_frames = cv2.VideoCapture(movie_name) + cur_frame = 0 + sum_frame = cap.get(cv2.CAP_PROP_FRAME_COUNT) + out_subdir = os.path.join(self.args.out_dir, '{}seq'.format(sequence_num)) + sequence_num += 1 + os.makedirs(out_subdir, exist_ok=True) + + while True: + ret, frame = cap.read() + if not ret: + print('{} / {}'.format(sum_frame, sum_frame)) + break + + if (cur_frame % self.args.pass_num) == 0: + frame = self.trim_endo_movie(frame) + # cv2.imshow("a", frame) + # cv2.waitKey() + frame = cv2.undistort(frame, self.intrinsics_row, self.dist_coeffs) + frame = cv2.resize(frame, self.resized_size, cv2.INTER_LANCZOS4) + # cv2.imshow("a", frame) + # cv2.waitKey() + frame = self.inpaint_endo_img(frame) + + cv2.imwrite(os.path.join(out_subdir, '{:08}.png'.format(file_name_num)), frame) + file_name_num += 1 + + if (cur_frame % 300) == 0: + print('{} / {}'.format(cur_frame, sum_frame)) + + cur_frame += 1 + + def trim_endo_movie(self, frame): + frame = frame[32:989, 323:1599, :] + frame = cv2.resize(frame, (self.calibrate_img_size[0], self.calibrate_img_size[1]), + interpolation=cv2.INTER_LANCZOS4) + return frame + + def inpaint_endo_img(self, frame): + YUV = cv2.cvtColor(frame, cv2.COLOR_BGR2YUV) + Y = YUV[:, :, 0] + _, mask = cv2.threshold(Y, self.args.thresh_num, 255, cv2.THRESH_BINARY) + kernel = np.ones((9, 9), np.uint8) + dilation = cv2.dilate(mask, kernel) + out = cv2.inpaint(frame, dilation, 30, cv2.INPAINT_NS) + return out diff --git a/__pycache__/EndoCalibrator.cpython-36.pyc b/__pycache__/EndoCalibrator.cpython-36.pyc index a6cfbdf..8c705ee 100644 --- a/__pycache__/EndoCalibrator.cpython-36.pyc +++ b/__pycache__/EndoCalibrator.cpython-36.pyc Binary files differ diff --git a/__pycache__/Movie2imgConverter.cpython-36.pyc b/__pycache__/Movie2imgConverter.cpython-36.pyc new file mode 100644 index 0000000..60f25b5 --- /dev/null +++ b/__pycache__/Movie2imgConverter.cpython-36.pyc Binary files differ diff --git a/main.py b/main.py index b2ed764..1b0bdf8 100644 --- a/main.py +++ b/main.py @@ -1,14 +1,15 @@ import argparse from EndoCalibrator import EndoCalibrator +from Movie2imgConverter import Movie2imgConverter parser = argparse.ArgumentParser() parser.add_argument("--mode", type=str, required=True, choices=["calibrate", "convert"], help="処理モード") -parser.add_argument("--imgs_dir", type=str, required=True, help="画像が入ってるディレクトリ") parser.add_argument("--resized_height", type=int, default=352, help="画像サイズの高さ（幅もだが，アップサンプリングの" "関係で2の乗数を多く含む方が良い)") parser.add_argument("--resized_width", type=int, default=480, help="画像サイズの幅") # カメラキャリブレーション用 +parser.add_argument("--imgs_dir", type=str, default="./tmp", help="チェス画像が入ってるディレクトリ,キャリブ時はrequired") parser.add_argument("--npy_outdir", type=str, default="./params", help="カメラキャリブレーション結果の出力先") parser.add_argument("--square_size", type=float, default=20.0, help='正方形の一辺のサイズ[mm]') parser.add_argument("--pattern_size_col", type=int, default=6, help="チェスの列数") @@ -19,8 +20,9 @@ parser.add_argument("--movie_dir", type=str, default="./tmp", help="入力動画のディレクトリ，映像変換時はrequired") parser.add_argument("--out_dir", type=str, default="./out", help="出力用のディレクトリ") parser.add_argument("--param_dir", type=str, default="./params", help="カメラキャリブレーションで取得したパラメータ先") -parser.add_argument("--num_of_pass", type=int, default=1, help="フレーム間の間") -parser.add_argument("--endo_movie_extend", type=str, default=".mov", help="入力動画の拡張子") +parser.add_argument("--pass_num", type=int, default=1, help="フレーム間の間(1で連続したフレーム)") +parser.add_argument("--endo_movie_extend", type=str, default=".mp4", help="入力動画の拡張子") +parser.add_argument("--thresh_num", type=int, default=180, help='inpaint用の閾値') args = parser.parse_args() @@ -30,3 +32,7 @@ if mode == 'calibrate': editor = EndoCalibrator(args) editor.calibrate_row_size(Is_example_showed=False) + + if mode == 'convert': + editor = Movie2imgConverter(args) + editor.convert_endo_movie2img() diff --git a/params/intrinsics_row.npy b/params/intrinsics_row.npy new file mode 100644 index 0000000..37715d3 --- /dev/null +++ b/params/intrinsics_row.npy Binary files differ