diff --git a/colorization/my_colorize.py b/colorization/my_colorize.py index dad05dcb8b1cf80b419bb9bf948281671341651c..d6c517929088289c7a9525a8f0ec6fbc2e7c0d98 100644 --- a/colorization/my_colorize.py +++ b/colorization/my_colorize.py @@ -94,6 +94,10 @@ def train_model(model, criterion, optimizer, scheduler, num_epochs=25): # load colorizer colorizer = siggraph17(pretrained=False) + +for name, param in colorizer.named_parameters(): + print(name) + postcard_dataset = ColDataset(csv_file='models/colored_postcards/akon_postcards_public_domain.csv', root_dir='models/colored_postcards/low_res_imgs') (train_set, test_set) = random_split(postcard_dataset, [round(0.7*len(postcard_dataset)), @@ -106,9 +110,10 @@ criterion = nn.SmoothL1Loss() optimizer = optim.Adam(colorizer.parameters(), lr=0.001) scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=7, gamma=0.1) -colorizer = train_model(colorizer, criterion, optimizer, scheduler) -PATH = './models/col_test.pth' -torch.save(colorizer.state_dict(), PATH) + +# colorizer = train_model(colorizer, criterion, optimizer, scheduler) +# PATH = './models/col_test.pth' +# torch.save(colorizer.state_dict(), PATH) # # img = load_img('models/colored_postcards/imgs/AK001_011.jpg') diff --git a/stitching/imgs/pano_1.jpg b/stitching/imgs/pano_1.jpg new file mode 100644 index 0000000000000000000000000000000000000000..21ff278f41d89eabfbf268620e6897b4d30771c2 Binary files /dev/null and b/stitching/imgs/pano_1.jpg differ diff --git a/stitching/imgs/pano_2.jpg b/stitching/imgs/pano_2.jpg new file mode 100644 index 0000000000000000000000000000000000000000..0a6237c05b514405c5e767ff3a78c1a83f8c9ae9 Binary files /dev/null and b/stitching/imgs/pano_2.jpg differ diff --git a/stitching/imgs/pano_3.jpg b/stitching/imgs/pano_3.jpg new file mode 100644 index 0000000000000000000000000000000000000000..6ba7ef56ce85b362e1590540dae18522ccedcf8c Binary files /dev/null and b/stitching/imgs/pano_3.jpg differ diff --git a/stitching/scan.py b/stitching/scan.py new file mode 100644 index 0000000000000000000000000000000000000000..516703ba838ca5d68271fed2929a44ce0a33126c --- /dev/null +++ b/stitching/scan.py @@ -0,0 +1,212 @@ +""" +Stitching sample (advanced) +=========================== +Show how to use Stitcher API from python. +Adapted from https://raw.githubusercontent.com/opencv/opencv/master/samples/python/stitching_detailed.py +""" +import cv2 as cv +import numpy as np + + +def resize_image(img, scale): + return cv.resize(src=img, dsize=None, fx=scale, fy=scale, interpolation=cv.INTER_LINEAR_EXACT) + + +def resize_mask(msk, target_msk): + return cv.resize(src=msk, dsize=(target_msk.shape[1], target_msk.shape[0]), fx=0, fy=0, + interpolation=cv.INTER_LINEAR_EXACT) + + +def find_features(imgs): + features = [] + finder = cv.ORB.create() + for img in imgs: + img_feat = cv.detail.computeImageFeatures2(finder, img) + features.append(img_feat) + return features + + +def match_features(features, match_conf=0.3): + matcher = cv.detail_AffineBestOf2NearestMatcher(False, False, match_conf) + matches = matcher.apply2(features) + matcher.collectGarbage() + return matches + + +def get_matches_graph(img_names, matches, conf_thresh=1.0): + return cv.detail.matchesGraphAsString(img_names, matches, conf_thresh) + + +def get_biggest_subset(img_names, imgs, features, matches, conf_thresh=1.0, match_conf=0.3): + indices = cv.detail.leaveBiggestComponent(features, matches, conf_thresh) + img_subset = [] + img_names_subset = [] + features_subset = [] + num = len(indices) + if num < 2: + print('Need more images') + exit() + for i in range(num): + img_subset.append(imgs[indices[i, 0]]) + img_names_subset.append(img_names[indices[i, 0]]) + features_subset.append(features[indices[i, 0]]) + matches_subset = match_features(features_subset, match_conf) + return indices, img_names_subset, img_subset, features_subset, matches_subset + + +def estimate_cameras(features, matches): + estimator = cv.detail_AffineBasedEstimator() + b, cameras = estimator.apply(features, matches, None) + if not b: + print('Homography estimation failed') + exit() + for cam in cameras: + cam.R = cam.R.astype(np.float32) + return cameras + + +def refine_cameras(features, matches, cameras, refine_mask): + adjuster = cv.detail_BundleAdjusterAffinePartial() + adjuster.setConfThresh(1) + mask = np.zeros((3, 3), np.uint8) + if refine_mask[0] == 'x': + mask[0, 0] = 1 + if refine_mask[1] == 'x': + mask[0, 1] = 1 + if refine_mask[2] == 'x': + mask[0, 2] = 1 + if refine_mask[3] == 'x': + mask[1, 1] = 1 + if refine_mask[4] == 'x': + mask[1, 2] = 1 + adjuster.setRefinementMask(mask) + b, refined_cams = adjuster.apply(features, matches, cameras) + if not b: + print('Camera parameters adjusting failed') + exit() + return refined_cams + + +def get_warped_image_scale(cameras): + focals = [cam.focal for cam in cameras] + focals.sort() + if len(focals) % 2 == 1: + warped_image_scale = focals[len(focals) // 2] + else: + warped_image_scale = (focals[len(focals) // 2] + focals[len(focals) // 2 - 1]) / 2 + return warped_image_scale + + +def get_scale(img, megapix): + return min(1.0, np.sqrt(megapix * 1e6 / (img.shape[0] * img.shape[1]))) + + +def get_warped_masks_and_images(imgs, cameras, warp_scale, seam_aspect): + corners = [] + masks = [] + masks_warped = [] + images_warped = [] + num_imgs = len(imgs) + sizes = get_sizes(imgs) + for i in range(0, num_imgs): + um = cv.UMat(255 * np.ones(sizes[i], np.uint8)) + masks.append(um) + warper = cv.PyRotationWarper('affine', warp_scale * seam_aspect) + for idx in range(num_imgs): + K = cameras[idx].K().astype(np.float32) + swa = seam_aspect + K[0, 0] *= swa + K[0, 2] *= swa + K[1, 1] *= swa + K[1, 2] *= swa + corner, image_wp = warper.warp(imgs[idx], K, cameras[idx].R, cv.INTER_LINEAR, cv.BORDER_REFLECT) + corners.append(corner) + images_warped.append(image_wp) + p, mask_wp = warper.warp(masks[idx], K, cameras[idx].R, cv.INTER_LINEAR, cv.BORDER_REFLECT) + masks_warped.append(mask_wp.get()) + return corners, images_warped, masks_warped + + +def compensate_exposure(corners, images_warped, masks_warped): + comp = cv.detail.ExposureCompensator_createDefault(cv.detail.ExposureCompensator_GAIN_BLOCKS) + comp.feed(corners, images_warped, masks_warped) + return comp + + +def find_seams(corners, imgs_w, msks_w): + imgs_w_f = [img.astype(np.float32) for img in imgs_w] + seam_finder = cv.detail_GraphCutSeamFinder('COST_COLOR') + seam_finder.find(imgs_w_f, corners, msks_w) + return msks_w + + +def resize_seam_masks(msks_w, orig_msks): + seam_msks = [] + for idx in range(len(msks_w)): + dilated_mask = cv.dilate(msks_w[idx], None) + seam_mask = resize_mask(dilated_mask, orig_msks[idx]) + mask_warped = cv.bitwise_and(seam_mask, orig_msks[idx]) + seam_msks.append(mask_warped) + return seam_msks + + +def get_sizes(imgs): + sizes = [] + for img in imgs: + x, y = img.shape[1], img.shape[0] + sizes.append((x, y)) + return sizes + + +def blend_images(corners, imgs, msks, msks_w, comp): + sizes = get_sizes(imgs) + blender = cv.detail_MultiBandBlender() + dst_sz = cv.detail.resultRoi(corners=corners, sizes=sizes) + blend_strength = 5 + blend_width = np.sqrt(dst_sz[2] * dst_sz[3]) * blend_strength / 100 + blender.setNumBands((np.log(blend_width) / np.log(2.) - 1.).astype(np.int32)) + blender.prepare(dst_sz) + for idx in range(len(imgs)): + comp.apply(idx, corners[idx], imgs[idx], msks[idx]) + img_s = imgs[idx].astype(np.int16) + msk_w = resize_seam_masks(msks[idx], msks_w[idx]) + blender.feed(cv.UMat(img_s), msk_w, corners[idx]) + result = None + result_mask = None + return blender.blend(result, result_mask) + + +def get_reg_data(imgs, img_names, reg_megapix=0.6): + reg_scale = get_scale(imgs[0], reg_megapix) + imgs_med_res = [resize_image(img, reg_scale) for img in imgs] + features = find_features(imgs_med_res) + matches = match_features(features) + indices, img_names_subset, img_subset, features_subset, matches_subset = get_biggest_subset( + img_names, imgs_med_res, features, matches, conf_thresh=0.3) + graph = get_matches_graph(img_names_subset, matches_subset, conf_thresh=0.3) + cameras = estimate_cameras(features_subset, matches_subset) + cameras = refine_cameras(features_subset, matches_subset, cameras, 'xxxxx') + return indices, graph, cameras + + +def compose_from_reg_data(imgs, indices, cameras, reg_scale, seam_scale): # TODO: fully implement composing step + warp_scale = get_warped_image_scale(cameras) + seam_aspect = seam_scale / reg_scale + corners, imgs_w, msks_w = get_warped_masks_and_images(imgs, cameras, warp_scale, seam_aspect) + return corners + + +def full_stitching_pipeline(imgs, img_names): + indices, graph, cameras = get_reg_data(imgs, img_names) + pano = compose_from_reg_data(imgs, indices, cameras) + return pano + + +if __name__ == '__main__': + img_names = [f'pano_{i}.jpg' for i in range(1, 4)] + imgs = [cv.imread(f'imgs/{name}') for name in img_names] + indices, graph, cameras = get_reg_data(imgs, img_names) + print(graph) + reg_scale = get_scale(imgs[0], megapix=0.6) + seam_scale = get_scale(imgs[0], megapix=0.1) + print(compose_from_reg_data(imgs, indices, cameras, reg_scale, seam_scale)) diff --git a/stitching/stitch_demo.py b/stitching/stitch_demo.py index f8ac3c2df72dd2057f4663b994084146491fcc82..c445ee2284db034ae61f5bc0ca967a1266486176 100644 --- a/stitching/stitch_demo.py +++ b/stitching/stitch_demo.py @@ -5,7 +5,7 @@ Show how to use Stitcher API from python. """ # Python 2/3 compatibility -from __future__ import print_function +# from __future__ import print_function import argparse from collections import OrderedDict @@ -383,6 +383,7 @@ def main(): if not b: print("Camera parameters adjusting failed.") exit() + print([cam.R for cam in cameras]) focals = [] for cam in cameras: focals.append(cam.focal) @@ -465,6 +466,7 @@ def main(): _img_size = (img.shape[1], img.shape[0]) K = cameras[idx].K().astype(np.float32) corner, image_warped = warper.warp(img, K, cameras[idx].R, cv.INTER_LINEAR, cv.BORDER_REFLECT) + print(corner) mask = 255 * np.ones((img.shape[0], img.shape[1]), np.uint8) p, mask_warped = warper.warp(mask, K, cameras[idx].R, cv.INTER_NEAREST, cv.BORDER_CONSTANT) compensator.apply(idx, corners[idx], image_warped, mask_warped) @@ -504,17 +506,17 @@ def main(): result_mask = None result, result_mask = blender.blend(result, result_mask) cv.imwrite(result_name, result) - zoom_x = 600.0 / result.shape[1] - dst = cv.normalize(src=result, dst=None, alpha=255., norm_type=cv.NORM_MINMAX, dtype=cv.CV_8U) - dst = cv.resize(dst, dsize=None, fx=zoom_x, fy=zoom_x) - cv.imshow(result_name, dst) - cv.waitKey() + # zoom_x = 600.0 / result.shape[1] + # dst = cv.normalize(src=result, dst=None, alpha=255., norm_type=cv.NORM_MINMAX, dtype=cv.CV_8U) + # dst = cv.resize(dst, dsize=None, fx=zoom_x, fy=zoom_x) + # cv.imshow(result_name, dst) + # cv.waitKey() - print(corners) + # print(corners) print("Done") if __name__ == '__main__': # print(__doc__) main() - cv.destroyAllWindows() \ No newline at end of file + cv.destroyAllWindows()