Skip to content

modosaic.normals.generators.impl.midas_d2n

modosaic.normals.generators.impl.midas_d2n

MidasD2N 

MidasD2N()

Bases: NormalsGenerator

MiDaS-based depth-to-normal generator.

Load the MiDaS model and transform.

Source code in modosaic/normals/generators/impl/midas_d2n.py 
20
21
22
23
24
def __init__(self):
    """Load the MiDaS model and transform."""
    super().__init__()
    self.model = self._load_model()
    self.transform = MidasD2N._get_model_transform()

generate 

generate(record)

Generate a normal field for one image record.

Source code in modosaic/normals/generators/impl/midas_d2n.py 
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
@override
def generate(self, record: ImageRecord) -> np.ndarray:
    """Generate a normal field for one image record."""
    logger.debug(f"Generating normals for record sample {record.sample_id}")
    pil_img = ImageService.bytes_to_pil(record.image_bytes)
    img = np.array(pil_img.convert("RGB"))
    inp = self.transform(img).to(self.device)

    pred = self.model(inp)
    pred = torch.nn.functional.interpolate(
        pred.unsqueeze(1),
        size=img.shape[:2],
        mode="bicubic",
        align_corners=False,
    ).squeeze(1)

    depth = pred[0].detach().float().cpu().numpy()
    lo, hi = np.percentile(depth, 2), np.percentile(depth, 98)
    d = (depth - lo) / (hi - lo + 1e-6)
    d = np.clip(d, 0.0, 1.0)

    normals = self._depth_to_normals(d)
    return normals