Function estimatemidas

src/depthmap_generation.py:455–499

(img, model, w, h, resize_mode, normalization, no_half, precision_is_autocast)

def estimatemidas(img, model, w, h, resize_mode, normalization, no_half, precision_is_autocast):
    import contextlib
    # init transform
    transform = Compose(
        [
            Resize(
                w,
                h,
                resize_target=None,
                keep_aspect_ratio=True,
                ensure_multiple_of=32,
                resize_method=resize_mode,
                image_interpolation_method=cv2.INTER_CUBIC,
            ),
            normalization,
            PrepareForNet(),
        ]
    )

    # transform input
    img_input = transform({"image": img})["image"]

    # compute
    precision_scope = torch.autocast if precision_is_autocast and depthmap_device == torch.device(
        "cuda") else contextlib.nullcontext
    with torch.no_grad(), precision_scope("cuda"):
        sample = torch.from_numpy(img_input).to(depthmap_device).unsqueeze(0)
        if depthmap_device == torch.device("cuda"):
            sample = sample.to(memory_format=torch.channels_last)
            if not no_half:
                sample = sample.half()
        prediction = model.forward(sample)
        prediction = (
            torch.nn.functional.interpolate(
                prediction.unsqueeze(1),
                size=img.shape[:2],
                mode="bicubic",
                align_corners=False,
            )
            .squeeze()
            .cpu()
            .numpy()
        )

    return prediction


# TODO: correct values for BOOST