Bases: NormalsGenerator
Omnidata surface-normal generator.
Load the Omnidata model and image transform.
Source code in modosaic/normals/generators/impl/omnidata.py
| def __init__(self):
"""Load the Omnidata model and image transform."""
super().__init__()
self.model = self._load_model()
self.transform = Omnidata._get_model_transform()
|
generate
Generate a normal field for one image record.
Source code in modosaic/normals/generators/impl/omnidata.py
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)
x = self.transform(pil_img.convert("RGB")).unsqueeze(0).to(self.device)
y = self.model(x) # (1,3,H,W)
y = y[0].detach().float().cpu().permute(1, 2, 0).numpy()
y = Omnidata._normalize_vec(y).astype(np.float32)
# Resize back to original size for visual comparison
width, height = pil_img.size
y_vis = Omnidata._normals_to_rgb(y).resize((width, height), resample=Image.BILINEAR)
y2 = (np.array(y_vis).astype(np.float32) / 255.0) * 2.0 - 1.0
normals = Omnidata._normalize_vec(y2).astype(np.float32)
return normals
|