Spaces:
Running
Running
| import sys | |
| from pathlib import Path | |
| from ..utils.base_model import BaseModel | |
| import torch | |
| from ..utils.base_model import BaseModel | |
| sold2_path = Path(__file__).parent / "../../third_party/SOLD2" | |
| sys.path.append(str(sold2_path)) | |
| from sold2.model.line_matcher import LineMatcher | |
| from sold2.misc.visualize_util import ( | |
| plot_images, | |
| plot_lines, | |
| plot_line_matches, | |
| plot_color_line_matches, | |
| plot_keypoints, | |
| ) | |
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
| class SOLD2(BaseModel): | |
| default_conf = { | |
| "weights": "sold2_wireframe.tar", | |
| "match_threshold": 0.2, | |
| "checkpoint_dir": sold2_path / "pretrained", | |
| "detect_thresh": 0.25, | |
| "multiscale": False, | |
| "valid_thresh": 1e-3, | |
| "num_blocks": 20, | |
| "overlap_ratio": 0.5, | |
| } | |
| required_inputs = [ | |
| "image0", | |
| "image1", | |
| ] | |
| # Initialize the line matcher | |
| def _init(self, conf): | |
| checkpoint_path = conf["checkpoint_dir"] / conf["weights"] | |
| mode = "dynamic" # 'dynamic' or 'static' | |
| match_config = { | |
| "model_cfg": { | |
| "model_name": "lcnn_simple", | |
| "model_architecture": "simple", | |
| # Backbone related config | |
| "backbone": "lcnn", | |
| "backbone_cfg": { | |
| "input_channel": 1, # Use RGB images or grayscale images. | |
| "depth": 4, | |
| "num_stacks": 2, | |
| "num_blocks": 1, | |
| "num_classes": 5, | |
| }, | |
| # Junction decoder related config | |
| "junction_decoder": "superpoint_decoder", | |
| "junc_decoder_cfg": {}, | |
| # Heatmap decoder related config | |
| "heatmap_decoder": "pixel_shuffle", | |
| "heatmap_decoder_cfg": {}, | |
| # Descriptor decoder related config | |
| "descriptor_decoder": "superpoint_descriptor", | |
| "descriptor_decoder_cfg": {}, | |
| # Shared configurations | |
| "grid_size": 8, | |
| "keep_border_valid": True, | |
| # Threshold of junction detection | |
| "detection_thresh": 0.0153846, # 1/65 | |
| "max_num_junctions": 300, | |
| # Threshold of heatmap detection | |
| "prob_thresh": 0.5, | |
| # Weighting related parameters | |
| "weighting_policy": mode, | |
| # [Heatmap loss] | |
| "w_heatmap": 0.0, | |
| "w_heatmap_class": 1, | |
| "heatmap_loss_func": "cross_entropy", | |
| "heatmap_loss_cfg": {"policy": mode}, | |
| # [Heatmap consistency loss] | |
| # [Junction loss] | |
| "w_junc": 0.0, | |
| "junction_loss_func": "superpoint", | |
| "junction_loss_cfg": {"policy": mode}, | |
| # [Descriptor loss] | |
| "w_desc": 0.0, | |
| "descriptor_loss_func": "regular_sampling", | |
| "descriptor_loss_cfg": { | |
| "dist_threshold": 8, | |
| "grid_size": 4, | |
| "margin": 1, | |
| "policy": mode, | |
| }, | |
| }, | |
| "line_detector_cfg": { | |
| "detect_thresh": 0.25, # depending on your images, you might need to tune this parameter | |
| "num_samples": 64, | |
| "sampling_method": "local_max", | |
| "inlier_thresh": 0.9, | |
| "use_candidate_suppression": True, | |
| "nms_dist_tolerance": 3.0, | |
| "use_heatmap_refinement": True, | |
| "heatmap_refine_cfg": { | |
| "mode": "local", | |
| "ratio": 0.2, | |
| "valid_thresh": 1e-3, | |
| "num_blocks": 20, | |
| "overlap_ratio": 0.5, | |
| }, | |
| }, | |
| "multiscale": False, | |
| "line_matcher_cfg": { | |
| "cross_check": True, | |
| "num_samples": 5, | |
| "min_dist_pts": 8, | |
| "top_k_candidates": 10, | |
| "grid_size": 4, | |
| }, | |
| } | |
| self.net = LineMatcher( | |
| match_config["model_cfg"], | |
| checkpoint_path, | |
| device, | |
| match_config["line_detector_cfg"], | |
| match_config["line_matcher_cfg"], | |
| match_config["multiscale"], | |
| ) | |
| def _forward(self, data): | |
| img0 = data["image0"] | |
| img1 = data["image1"] | |
| pred = self.net([img0, img1]) | |
| line_seg1 = pred["line_segments"][0] | |
| line_seg2 = pred["line_segments"][1] | |
| matches = pred["matches"] | |
| valid_matches = matches != -1 | |
| match_indices = matches[valid_matches] | |
| matched_lines1 = line_seg1[valid_matches][:, :, ::-1] | |
| matched_lines2 = line_seg2[match_indices][:, :, ::-1] | |
| pred["raw_lines0"], pred["raw_lines1"] = line_seg1, line_seg2 | |
| pred["lines0"], pred["lines1"] = matched_lines1, matched_lines2 | |
| pred = {**pred, **data} | |
| return pred | |