class_docstr: Module, which detects and describe line segments in an image.
  
  This is based on the original code from the paper "SOLD²: Self-supervised
  Occlusion-aware Line Detector and Descriptor". See :cite:`SOLD22021` for more details.
  
  Args:
      config: Dict specifying parameters. None will load the default parameters,
          which are tuned for images in the range 400~800 px.
      pretrained: If True, download and set pretrained weights to the model.
  
  Returns:
      The raw junction and line heatmaps, the semi-dense descriptor map,
      as well as the list of detected line segments (ij coordinates convention).
  
  Example:
      >>> images = torch.rand(2, 1, 512, 512)
      >>> sold2 = SOLD2()
      >>> outputs = sold2(images)
      >>> line_seg1 = outputs["line_segments"][0]
      >>> line_seg2 = outputs["line_segments"][1]
      >>> desc1 = outputs["dense_desc"][0]
      >>> desc2 = outputs["dense_desc"][1]
      >>> matches = sold2.match(line_seg1, line_seg2, desc1[None], desc2[None])
class_docstr_tok: ["Module",",","which","detects","and","describe","line","segments","in","an","image",".","This","is","based","on","the","original","code","from","the","paper","``","SOLD\u00b2",":","Self-supervised","Occlusion-aware","Line","Detector","and","Descriptor","''",".","See",":","cite",":","`","SOLD22021","`","for","more","details",".","Args",":","config",":","Dict","specifying","parameters",".","None","will","load","the","default","parameters",",","which","are","tuned","for","images","in","the","range","400~800","px",".","pretrained",":","If","True",",","download","and","set","pretrained","weights","to","the","model",".","Returns",":","The","raw","junction","and","line","heatmaps",",","the","semi-dense","descriptor","map",",","as","well","as","the","list","of","detected","line","segments","(","ij","coordinates","convention",")",".","Example",":",">",">",">","images","=","torch.rand","(","2",",","1",",","512",",","512",")",">",">",">","sold2","=","SOLD2","(",")",">",">",">","outputs","=","sold2","(","images",")",">",">",">","line_seg1","=","outputs","[","``","line_segments","''","]","[","0","]",">",">",">","line_seg2","=","outputs","[","``","line_segments","''","]","[","1","]",">",">",">","desc1","=","outputs","[","``","dense_desc","''","]","[","0","]",">",">",">","desc2","=","outputs","[","``","dense_desc","''","]","[","1","]",">",">",">","matches","=","sold2.match","(","line_seg1",",","line_seg2",",","desc1","[","None","]",",","desc2","[","None","]",")"]
docstr: Args:
      img: batched images with shape :math:`(B, 1, H, W)`.
  
  Return:
      - ``line_segments``: list of N line segments in each of the B images :math:`List[(N, 2, 2)]`.
      - ``junction_heatmap``: raw junction heatmap of shape :math:`(B, H, W)`.
      - ``line_heatmap``: raw line heatmap of shape :math:`(B, H, W)`.
      - ``dense_desc``: the semi-dense descriptor map of shape :math:`(B, 128, H/4, W/4)`.
docstr_tok: ["Args",":","img",":","batched","images","with","shape",":","math",":","`","(","B",",","1",",","H",",","W",")","`",".","Return",":","-","``","line_segments","``",":","list","of","N","line","segments","in","each","of","the","B","images",":","math",":","`","List","[","(","N",",","2",",","2",")","]","`",".","-","``","junction_heatmap","``",":","raw","junction","heatmap","of","shape",":","math",":","`","(","B",",","H",",","W",")","`",".","-","``","line_heatmap","``",":","raw","line","heatmap","of","shape",":","math",":","`","(","B",",","H",",","W",")","`",".","-","``","dense_desc","``",":","the","semi-dense","descriptor","map","of","shape",":","math",":","`","(","B",",","128",",","H\/4",",","W\/4",")","`","."]
code: def forward(self, img: Tensor) -> Dict[str, Any]:
      """
      Args:
          img: batched images with shape :math:`(B, 1, H, W)`.
  
      Return:
          - ``line_segments``: list of N line segments in each of the B images :math:`List[(N, 2, 2)]`.
          - ``junction_heatmap``: raw junction heatmap of shape :math:`(B, H, W)`.
          - ``line_heatmap``: raw line heatmap of shape :math:`(B, H, W)`.
          - ``dense_desc``: the semi-dense descriptor map of shape :math:`(B, 128, H/4, W/4)`.
      """
      KORNIA_CHECK_SHAPE(img, ["B", "1", "H", "W"])
      outputs = {}
  
      # Forward pass of the CNN backbone
      net_outputs = self.model(img)
      outputs["junction_heatmap"] = net_outputs["junctions"]
      outputs["line_heatmap"] = net_outputs["heatmap"]
      outputs["dense_desc"] = net_outputs["descriptors"]
  
      # Loop through all images
      lines = []
      for junc_prob, heatmap in zip(
          net_outputs["junctions"], net_outputs["heatmap"]
      ):
          # Get the junctions
          junctions = prob_to_junctions(
              junc_prob,
              self.grid_size,
              self.junc_detect_thresh,
              self.max_num_junctions,
          )
  
          # Run the line detector
          line_map, junctions, _ = self.line_detector.detect(
              junctions, heatmap
          )
          lines.append(line_map_to_segments(junctions, line_map))
      outputs["line_segments"] = lines
  
      return outputs
  
imports: from typing import Any, Dict, Optional, Tuple
  import torch
  import torch.nn.functional
  from kornia.core import Module, Tensor, concatenate, pad, stack
  from kornia.core.check import KORNIA_CHECK_SHAPE
  from kornia.geometry.conversions import normalize_pixel_coordinates
  from kornia.utils import map_location_to_cpu
  from .backbones import SOLD2Net
  from .sold2_detector import LineSegmentDetectionModule, line_map_to_segments, prob_to_junctions
label_desc: Use image node_id 2 for calling the SOLD2 obj's underlying member method code with example usage: obj.forward(img) and returns: outputs