jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def forward(self, lqs):
        """Forward function for BasicVSR++.

        Args:
            lqs (tensor): Input low quality (LQ) sequence with
                shape (n, t, c, h, w).

        Returns:
            Tensor: Output HR sequence with shape (n, t, c, 4h, 4w).
        """

        n, t, c, h, w = lqs.size()

        # whether to cache the features in CPU (no effect if using CPU)
        if t > self.cpu_cache_length and lqs.is_cuda:
            self.cpu_cache = True
        else:
            self.cpu_cache = False

        if self.is_low_res_input:
            lqs_downsample = lqs.clone()
        else:
            lqs_downsample = t(
                lqs.view(-1, c, h, w), scale_factor=0.25,
                mode='bicubic').view(n, t, c, h // 4, w // 4)

        # check whether the input is an extended sequence
        self.check_if_mirror_extended(lqs)

        feats = {}
        # compute spatial features
        if self.cpu_cache:
            feats['spatial'] = []
            for i in range(0, t):
                feat = self.feat_extract(lqs[:, i, :, :, :]).cpu()
                feats['spatial'].append(feat)
                torch.cuda.empty_cache()
        else:
            feats_ = self.feat_extract(lqs.view(-1, c, h, w))
            h, w = feats_.shape[2:]
            feats_ = feats_.view(n, t, -1, h, w)
            feats['spatial'] = [feats_[:, i, :, :, :] for i in range(0, t)]

        # compute optical flow using the low-res inputs
        assert lqs_downsample.size(3) >= 64 and lqs_downsample.size(4) >= 64, (
            'The height and width of low-res inputs must be at least 64, '
            f'but got {h} and {w}.')
        flows_forward, flows_backward = self.compute_flow(lqs_downsample)

        # feature propagation
        for iter_ in [1, 2]:
            for direction in ['backward', 'forward']:
                module = f'{direction}_{iter_}'

                feats[module] = []

                if direction == 'backward':
                    flows = flows_backward
                elif flows_forward is not None:
                    flows = flows_forward
                else:
                    flows = flows_backward.flip(1)

                feats = self.propagate(feats, flows, module)
                if self.cpu_cache:
                    del flows
                    torch.cuda.empty_cache()

        return self.upsample(lqs, feats)
//  [python]
def forward(self, lqs):
        """Forward function for BasicVSR++.

        Args:
            lqs (tensor): Input low quality (LQ) sequence with
                shape (n, t, c, h, w).

        Returns:
            Tensor: Output HR sequence with shape (n, t, c, 4h, 4w).
        """

        n, t, c, h, w = lqs.size()

        # whether to cache the features in CPU (no effect if using CPU)
        if t > self.cpu_cache_length and lqs.is_cuda:
            self.cpu_cache = True
        else:
            self.cpu_cache = False

        if self.is_low_res_input:
            lqs_downsample = lqs.clone()
        else:
            lqs_downsample = F.interpolate(
                lqs.view(-1, c, h, w), scale_factor=0.25,
                mode='bicubic').view(n, t, c, h // 4, w // 4)

        # check whether the input is an extended sequence
        self.check_if_mirror_extended(lqs)

        feats = {}
        # compute spatial features
        if self.cpu_cache:
            feats['spatial'] = []
            for i in range(0, t):
                feat = self.feat_extract(lqs[:, i, :, :, :]).cpu()
                feats['spatial'].append(feat)
                torch.cuda.empty_cache()
        else:
            feats_ = self.feat_extract(lqs.view(-1, c, h, w))
            h, w = feats_.shape[2:]
            feats_ = feats_.view(n, t, -1, h, w)
            feats['spatial'] = [feats_[:, i, :, :, :] for i in range(0, t)]

        # compute optical flow using the low-res inputs
        assert lqs_downsample.size(3) >= 64 and lqs_downsample.size(4) >= 64, (
            'The height and width of low-res inputs must be at least 64, '
            f'but got {h} and {w}.')
        flows_forward, flows_backward = self.compute_flow(lqs_downsample)

        # feature propagation
        for iter_ in [1, 2]:
            for direction in ['backward', 'forward']:
                module = f'{direction}_{iter_}'

                feats[module] = []

                if direction == 'backward':
                    flows = flows_backward
                elif flows_forward is not None:
                    flows = flows_forward
                else:
                    flows = flows_backward.flip(1)

                feats = self.propagate(feats, flows, module)
                if self.cpu_cache:
                    del flows
                    torch.cuda.empty_cache()

        return self.upsample(lqs, feats)
//  [python]
def forward(self, top, lateral):
        if lateral.shape[2] != top.shape[2]:
            h, w = lateral.size(2), lateral.size(3)
            top = F.upsample(input=top, size=(h, w), mode='bilinear')
        out = torch.cat((lateral, top), dim=1)
        out = self.relu(self.bn1(self.conv1(out)))
        out = self.relu(self.bn2(self.conv2(out)))
        return out
//  [python]
def forward(self, top, lateral):
        if lateral.shape[2] != top.shape[2]:
            h, w = lateral.size(2), lateral.size(3)
            top = F.interpolate(input=top, size=(h, w), mode='bilinear',align_corners=True)
        out = torch.cat((lateral, top), dim=1)
        out = self.relu(self.bn1(self.conv1(out)))
        out = self.relu(self.bn2(self.conv2(out)))
        return out
//  [python]
def get_input(
        self, batch, k, cond_key=None, bs=None, return_first_stage_outputs=False
    ):
        # note: restricted to non-trainable encoders currently
        assert (
            not self.cond_stage_trainable
        ), "trainable cond stages not yet supported for inpainting"
        z, c, x, xrec, xc = super().get_input(
            batch,
            self.first_stage_key,
            return_first_stage_outputs=True,
            force_c_encode=True,
            return_original_cond=True,
            bs=bs,
        )

        assert exists(self.concat_keys)
        c_cat = list()
        for ck in self.concat_keys:
            cc = (
                rearrange(batch[ck], "b h w c -> b c h w")
                .to(memory_format=torch.contiguous_format)
                .float()
            )
            if bs is not None:
                cc = cc[:bs]
                cc = cc.to(self.device)
            bchw = z.shape
            if ck != self.masked_image_key:
                cc = t.interpolate(cc, size=bchw[-2:])
            else:
                cc = self.get_first_stage_encoding(self.encode_first_stage(cc))
            c_cat.append(cc)
        c_cat = torch.cat(c_cat, dim=1)
        all_conds = {"c_concat": [c_cat], "c_crossattn": [c]}
        if return_first_stage_outputs:
            return z, all_conds, x, xrec, xc
        return z, all_conds
//  [python]
def get_input(
        self, batch, k, cond_key=None, bs=None, return_first_stage_outputs=False
    ):
        # note: restricted to non-trainable encoders currently
        assert (
            not self.cond_stage_trainable
        ), "trainable cond stages not yet supported for inpainting"
        z, c, x, xrec, xc = super().get_input(
            batch,
            self.first_stage_key,
            return_first_stage_outputs=True,
            force_c_encode=True,
            return_original_cond=True,
            bs=bs,
        )

        assert exists(self.concat_keys)
        c_cat = list()
        for ck in self.concat_keys:
            cc = (
                rearrange(batch[ck], "b h w c -> b c h w")
                .to(memory_format=torch.contiguous_format)
                .float()
            )
            if bs is not None:
                cc = cc[:bs]
                cc = cc.to(self.device)
            bchw = z.shape
            if ck != self.masked_image_key:
                cc = torch.nn.functional.interpolate(cc, size=bchw[-2:])
            else:
                cc = self.get_first_stage_encoding(self.encode_first_stage(cc))
            c_cat.append(cc)
        c_cat = torch.cat(c_cat, dim=1)
        all_conds = {"c_concat": [c_cat], "c_crossattn": [c]}
        if return_first_stage_outputs:
            return z, all_conds, x, xrec, xc
        return z, all_conds
//  [python]
def forward(self, x, output_shape=None):
        assert x.shape[1] == self.channels
        if self.dims == 3:
            shape = [x.shape[2], x.shape[3] * 2, x.shape[4] * 2]
            if output_shape is not None:
                shape[1] = output_shape[3]
                shape[2] = output_shape[4]
        else:
            shape = [x.shape[2] * 2, x.shape[3] * 2]
            if output_shape is not None:
                shape[0] = output_shape[2]
                shape[1] = output_shape[3]

        x = F.upsample(x, size=shape, mode="nearest")
        if self.use_conv:
            x = self.conv(x)
        return x
//  [python]
def forward(self, x, output_shape=None):
        assert x.shape[1] == self.channels
        if self.dims == 3:
            shape = [x.shape[2], x.shape[3] * 2, x.shape[4] * 2]
            if output_shape is not None:
                shape[1] = output_shape[3]
                shape[2] = output_shape[4]
        else:
            shape = [x.shape[2] * 2, x.shape[3] * 2]
            if output_shape is not None:
                shape[0] = output_shape[2]
                shape[1] = output_shape[3]

        x = F.interpolate(x, size=shape, mode="nearest")
        if self.use_conv:
            x = self.conv(x)
        return x
//  [python]
def forward(self, x, skip=None, upscale=True):
        """
        Applies the decoder block to the input tensor.

        Args:
            x (torch tensor): Input tensor.
            skip (torch tensor, optional): Skip connection tensor. Defaults to None.
            upscale (bool, optional): Whether to perform upscaling of the input tensor. Defaults to True.

        Returns:
            torch tensor: Output tensor.
        """
        if self.upscale and upscale:
            if self.use_pixel_shuffle:
                x = self.pixel_shuffle(x)
            else:
                x = t.upsample(x, scale_factor=2, mode="nearest")
        if skip is not None:
            x = torch.cat([x, skip], dim=1)
            x = self.attention1(x)
        x = self.conv1(x)
        x = self.conv2(x)
        x = self.attention2(x)
        return x
//  [python]
def forward(self, x, skip=None, upscale=True):
        """
        Applies the decoder block to the input tensor.

        Args:
            x (torch tensor): Input tensor.
            skip (torch tensor, optional): Skip connection tensor. Defaults to None.
            upscale (bool, optional): Whether to perform upscaling of the input tensor. Defaults to True.

        Returns:
            torch tensor: Output tensor.
        """
        if self.upscale and upscale:
            if self.use_pixel_shuffle:
                x = self.pixel_shuffle(x)
            else:
                x = F.interpolate(x, scale_factor=2, mode="nearest")
        if skip is not None:
            x = torch.cat([x, skip], dim=1)
            x = self.attention1(x)
        x = self.conv1(x)
        x = self.conv2(x)
        x = self.attention2(x)
        return x
//  [python]
def get_input(
        self, batch, k, cond_key=None, bs=None, return_first_stage_outputs=False
    ):
        # note: restricted to non-trainable encoders currently
        assert (
            not self.cond_stage_trainable
        ), "trainable cond stages not yet supported for depth2img"
        z, c, x, xrec, xc = super().get_input(
            batch,
            self.first_stage_key,
            return_first_stage_outputs=True,
            force_c_encode=True,
            return_original_cond=True,
            bs=bs,
        )

        assert exists(self.concat_keys)
        assert len(self.concat_keys) == 1
        c_cat = list()
        for ck in self.concat_keys:
            cc = batch[ck]
            if bs is not None:
                cc = cc[:bs]
                cc = cc.to(self.device)
            cc = self.depth_model(cc)
            cc = t.interpolate(
                cc,
                size=z.shape[2:],
                mode="bicubic",
                align_corners=False,
            )

            depth_min, depth_max = torch.amin(
                cc, dim=[1, 2, 3], keepdim=True
            ), torch.amax(cc, dim=[1, 2, 3], keepdim=True)
            cc = 2.0 * (cc - depth_min) / (depth_max - depth_min + 0.001) - 1.0
            c_cat.append(cc)
        c_cat = torch.cat(c_cat, dim=1)
        all_conds = {"c_concat": [c_cat], "c_crossattn": [c]}
        if return_first_stage_outputs:
            return z, all_conds, x, xrec, xc
        return z, all_conds
//  [python]
def get_input(
        self, batch, k, cond_key=None, bs=None, return_first_stage_outputs=False
    ):
        # note: restricted to non-trainable encoders currently
        assert (
            not self.cond_stage_trainable
        ), "trainable cond stages not yet supported for depth2img"
        z, c, x, xrec, xc = super().get_input(
            batch,
            self.first_stage_key,
            return_first_stage_outputs=True,
            force_c_encode=True,
            return_original_cond=True,
            bs=bs,
        )

        assert exists(self.concat_keys)
        assert len(self.concat_keys) == 1
        c_cat = list()
        for ck in self.concat_keys:
            cc = batch[ck]
            if bs is not None:
                cc = cc[:bs]
                cc = cc.to(self.device)
            cc = self.depth_model(cc)
            cc = torch.nn.functional.interpolate(
                cc,
                size=z.shape[2:],
                mode="bicubic",
                align_corners=False,
            )

            depth_min, depth_max = torch.amin(
                cc, dim=[1, 2, 3], keepdim=True
            ), torch.amax(cc, dim=[1, 2, 3], keepdim=True)
            cc = 2.0 * (cc - depth_min) / (depth_max - depth_min + 0.001) - 1.0
            c_cat.append(cc)
        c_cat = torch.cat(c_cat, dim=1)
        all_conds = {"c_concat": [c_cat], "c_crossattn": [c]}
        if return_first_stage_outputs:
            return z, all_conds, x, xrec, xc
        return z, all_conds
//  [python]
def compute_depth_losses_Hab(self, inputs, outputs, losses):
        """
        Compute depth metrics, to allow monitoring during training
        This isn't particularly accurate as it averages over the entire batch,
        so is only used to give an indication of validation performance
        """
        depth_pred = outputs[("depth", 0, 0)]
        depth_pred = torch.clamp(t.interpolate(
            depth_pred, [512, 512], mode="bilinear", align_corners=False), 1e-3, 10)
        depth_pred = depth_pred.detach()

        depth_gt = inputs["depth_gt"]
        mask = torch.logical_and(depth_gt>0.01, depth_gt<=10.0)

        depth_gt = depth_gt[mask]
        depth_pred = depth_pred[mask]
        depth_pred *= torch.median(depth_gt) / torch.median(depth_pred)
        depth_pred = torch.clamp(depth_pred, min=1e-3, max=10)
        depth_errors = compute_depth_errors(depth_gt, depth_pred)

        for i, metric in enumerate(self.depth_metric_names):
            losses[metric] = np.array(depth_errors[i].cpu())
//  [python]
def compute_depth_losses_Hab(self, inputs, outputs, losses):
        """
        Compute depth metrics, to allow monitoring during training
        This isn't particularly accurate as it averages over the entire batch,
        so is only used to give an indication of validation performance
        """
        depth_pred = outputs[("depth", 0, 0)]
        depth_pred = torch.clamp(F.interpolate(
            depth_pred, [512, 512], mode="bilinear", align_corners=False), 1e-3, 10)
        depth_pred = depth_pred.detach()

        depth_gt = inputs["depth_gt"]
        mask = torch.logical_and(depth_gt>0.01, depth_gt<=10.0)

        depth_gt = depth_gt[mask]
        depth_pred = depth_pred[mask]
        depth_pred *= torch.median(depth_gt) / torch.median(depth_pred)
        depth_pred = torch.clamp(depth_pred, min=1e-3, max=10)
        depth_errors = compute_depth_errors(depth_gt, depth_pred)

        for i, metric in enumerate(self.depth_metric_names):
            losses[metric] = np.array(depth_errors[i].cpu())
//  [python]
def generate_images_pred(self, inputs, outputs):
        """Generate the warped (reprojected) color images for a minibatch.
        Generated images are saved into the `outputs` dictionary.
        """
        for scale in self.opt.scales:
            disp = outputs[("out", scale)]
            disp = t.nn.functional.interpolate(
                disp, [self.opt.height, self.opt.width], mode="bilinear", align_corners=False)
            source_scale = 0

            depth = output_to_depth(disp, self.opt.min_depth, self.opt.max_depth)

            outputs[("depth", 0, scale)] = depth

            for i, frame_id in enumerate(self.opt.frame_ids[1:]):

                if frame_id == "s":
                    T = inputs["stereo_T"]
                else:
                    T = outputs[("cam_T_cam", 0, frame_id)]

                cam_points = self.backproject_depth[source_scale](
                    depth, inputs[("inv_K", source_scale)])
                pix_coords = self.project_3d[source_scale](
                    cam_points, inputs[("K", source_scale)], T)

                outputs[("sample", frame_id, scale)] = pix_coords

                outputs[("color", frame_id, scale)] = t.nn.functional.grid_sample(
                    inputs[("color", frame_id, source_scale)],
                    outputs[("sample", frame_id, scale)],
                    padding_mode="border")

                # auto-masking
                outputs[("color_identity", frame_id, scale)] = \
                    inputs[("color", frame_id, source_scale)]
//  [python]
def generate_images_pred(self, inputs, outputs):
        """Generate the warped (reprojected) color images for a minibatch.
        Generated images are saved into the `outputs` dictionary.
        """
        for scale in self.opt.scales:
            disp = outputs[("out", scale)]
            disp = F.interpolate(
                disp, [self.opt.height, self.opt.width], mode="bilinear", align_corners=False)
            source_scale = 0

            depth = output_to_depth(disp, self.opt.min_depth, self.opt.max_depth)

            outputs[("depth", 0, scale)] = depth

            for i, frame_id in enumerate(self.opt.frame_ids[1:]):

                if frame_id == "s":
                    T = inputs["stereo_T"]
                else:
                    T = outputs[("cam_T_cam", 0, frame_id)]

                cam_points = self.backproject_depth[source_scale](
                    depth, inputs[("inv_K", source_scale)])
                pix_coords = self.project_3d[source_scale](
                    cam_points, inputs[("K", source_scale)], T)

                outputs[("sample", frame_id, scale)] = pix_coords

                outputs[("color", frame_id, scale)] = F.grid_sample(
                    inputs[("color", frame_id, source_scale)],
                    outputs[("sample", frame_id, scale)],
                    padding_mode="border")

                # auto-masking
                outputs[("color_identity", frame_id, scale)] = \
                    inputs[("color", frame_id, source_scale)]
//  [python]
def generate_images_pred_multi(self, inputs, outputs):
        for scale in self.opt.scales:
            disp = outputs[("out", scale)]
            disp = t.nn.functional.interpolate(
                disp, [self.opt.height, self.opt.width], mode="bilinear", align_corners=False)
            depth = output_to_depth(disp, self.opt.min_depth, self.opt.max_depth)
            outputs[("depth", 0, scale)] = depth
//  [python]
def generate_images_pred_multi(self, inputs, outputs):
        for scale in self.opt.scales:
            disp = outputs[("out", scale)]
            disp = F.interpolate(
                disp, [self.opt.height, self.opt.width], mode="bilinear", align_corners=False)
            depth = output_to_depth(disp, self.opt.min_depth, self.opt.max_depth)
            outputs[("depth", 0, scale)] = depth
//  [python]
def compute_depth_errors_VA(self, inputs, outputs, losses):
        """
        compute depth errors on VA
        """
        depth_pred = outputs[("depth", 0, 0)]
        depth_pred = t.interpolate(depth_pred, [640, 640], mode="bilinear", align_corners=True)
        depth_pred = depth_pred.detach()

        depth_gt = inputs["depth_gt"]
        mask = torch.logical_and(depth_gt > 0.01, depth_gt<=10.0)

        depth_gt = depth_gt[mask]
        depth_pred = depth_pred[mask]
        depth_pred *= torch.median(depth_gt) / torch.median(depth_pred)
        depth_pred = torch.clamp(depth_pred, min=1e-3, max=10.0)
        depth_errors = compute_depth_errors(depth_gt, depth_pred)

        if losses is None:
            losses = {}
        for i, metric in enumerate(self.depth_metric_names):
            losses[metric] = np.array(depth_errors[i].cpu())
//  [python]
def compute_depth_errors_VA(self, inputs, outputs, losses):
        """
        compute depth errors on VA
        """
        depth_pred = outputs[("depth", 0, 0)]
        depth_pred = F.interpolate(depth_pred, [640, 640], mode="bilinear", align_corners=True)
        depth_pred = depth_pred.detach()

        depth_gt = inputs["depth_gt"]
        mask = torch.logical_and(depth_gt > 0.01, depth_gt<=10.0)

        depth_gt = depth_gt[mask]
        depth_pred = depth_pred[mask]
        depth_pred *= torch.median(depth_gt) / torch.median(depth_pred)
        depth_pred = torch.clamp(depth_pred, min=1e-3, max=10.0)
        depth_errors = compute_depth_errors(depth_gt, depth_pred)

        if losses is None:
            losses = {}
        for i, metric in enumerate(self.depth_metric_names):
            losses[metric] = np.array(depth_errors[i].cpu())
//  [python]
def compute_depth_errors_NYUv2(self, inputs, outputs, losses):
        """
        compute depth errors on NYUv2
        """
        depth_pred = outputs[("depth", 0, 0)]
        depth_pred = t.interpolate(depth_pred, [448, 608], mode="bilinear", align_corners=True)
        depth_pred = depth_pred.detach()

        depth_gt = inputs["depth_gt"]
        mask = torch.logical_and(depth_gt > 0.01, depth_gt<=10.0)

        depth_gt = depth_gt[mask]
        depth_pred = depth_pred[mask]
        depth_pred *= torch.median(depth_gt) / torch.median(depth_pred)
        depth_pred = torch.clamp(depth_pred, min=1e-3, max=10.0)
        depth_errors = compute_depth_errors(depth_gt, depth_pred)

        if losses is None:
            losses = {}
        for i, metric in enumerate(self.depth_metric_names):
            losses[metric] = np.array(depth_errors[i].cpu())
//  [python]
def compute_depth_errors_NYUv2(self, inputs, outputs, losses):
        """
        compute depth errors on NYUv2
        """
        depth_pred = outputs[("depth", 0, 0)]
        depth_pred = F.interpolate(depth_pred, [448, 608], mode="bilinear", align_corners=True)
        depth_pred = depth_pred.detach()

        depth_gt = inputs["depth_gt"]
        mask = torch.logical_and(depth_gt > 0.01, depth_gt<=10.0)

        depth_gt = depth_gt[mask]
        depth_pred = depth_pred[mask]
        depth_pred *= torch.median(depth_gt) / torch.median(depth_pred)
        depth_pred = torch.clamp(depth_pred, min=1e-3, max=10.0)
        depth_errors = compute_depth_errors(depth_gt, depth_pred)

        if losses is None:
            losses = {}
        for i, metric in enumerate(self.depth_metric_names):
            losses[metric] = np.array(depth_errors[i].cpu())
//  [python]
def forward(self, x):
        out_height, out_width = x.size()[2:]
        x = self.context_embedding(x)
        return F.upsample(x, size=(out_height, out_width), mode="nearest")
//  [python]
def forward(self, x):
        out_height, out_width = x.size()[2:]
        x = self.context_embedding(x)
        return F.interpolate(x, size=(out_height, out_width), mode="nearest")
//  [python]
def eval_save(self):
        """
        save prediction for a minibatch
        """
        self.set_eval()
        try:
            inputs = self.val_iter.__next__()
        except StopIteration:
            self.val_iter = iter(self.val_loader)
            inputs = self.val_iter.__next__()

        with torch.no_grad():
            inputs[("color_aug", 0, 0)] = inputs[("color_aug", 0, 0)].cuda()
            features = self.models["encoder"](inputs[("color_aug", 0, 0)]) #
            outputs = self.models["depth"](features)
            depth = output_to_depth(outputs[('out', 0)], self.opt.min_depth, self.opt.max_depth)
            outputs[("depth", 0, 0)] = depth

            sz = (640,640)
            store_path = 'results/VA'
            if not os.path.exists(store_path):
                os.makedirs(store_path)

            img = inputs[('color_aug',0, 0)]
            img = t.nn.functional.interpolate(img, sz, mode='bilinear', align_corners=True)

            img = img.cpu().numpy().squeeze()
            if img.shape[0] == self.opt.batch_size: # batch_size > 1
                img = np.transpose(img, (0,2,3,1))
            else: # batch_size == 1
                img = np.transpose(img, (1,2,0))

            depth = outputs[('depth', 0, 0)] * self.approx_factor #approximate alignment for visualization
            depth = t.nn.functional.interpolate(depth, sz, mode='bilinear', align_corners=True)
            depth = depth.cpu().numpy().squeeze()

            bsz = img.shape[0]

            if bsz != self.opt.batch_size: # save only one image
                imageio.imwrite(f'{store_path}/00_current.png', img)
                write_turbo_depth_metric(f'{store_path}/00_depth.png', depth, vmax=10.0)
            else: # forloop to save images
                for idx in range(bsz):
                    imageio.imwrite(f'{store_path}/{idx:02d}_current.png', img[idx])
                    write_turbo_depth_metric(f'{store_path}/{idx:02d}_depth.png', depth[idx], vmax=10.0)

            del inputs, outputs
//  [python]
def eval_save(self):
        """
        save prediction for a minibatch
        """
        self.set_eval()
        try:
            inputs = self.val_iter.__next__()
        except StopIteration:
            self.val_iter = iter(self.val_loader)
            inputs = self.val_iter.__next__()

        with torch.no_grad():
            inputs[("color_aug", 0, 0)] = inputs[("color_aug", 0, 0)].cuda()
            features = self.models["encoder"](inputs[("color_aug", 0, 0)]) #
            outputs = self.models["depth"](features)
            depth = output_to_depth(outputs[('out', 0)], self.opt.min_depth, self.opt.max_depth)
            outputs[("depth", 0, 0)] = depth

            sz = (640,640)
            store_path = 'results/VA'
            if not os.path.exists(store_path):
                os.makedirs(store_path)

            img = inputs[('color_aug',0, 0)]
            img = F.interpolate(img, sz, mode='bilinear', align_corners=True)

            img = img.cpu().numpy().squeeze()
            if img.shape[0] == self.opt.batch_size: # batch_size > 1
                img = np.transpose(img, (0,2,3,1))
            else: # batch_size == 1
                img = np.transpose(img, (1,2,0))

            depth = outputs[('depth', 0, 0)] * self.approx_factor #approximate alignment for visualization
            depth = F.interpolate(depth, sz, mode='bilinear', align_corners=True)
            depth = depth.cpu().numpy().squeeze()

            bsz = img.shape[0]

            if bsz != self.opt.batch_size: # save only one image
                imageio.imwrite(f'{store_path}/00_current.png', img)
                write_turbo_depth_metric(f'{store_path}/00_depth.png', depth, vmax=10.0)
            else: # forloop to save images
                for idx in range(bsz):
                    imageio.imwrite(f'{store_path}/{idx:02d}_current.png', img[idx])
                    write_turbo_depth_metric(f'{store_path}/{idx:02d}_depth.png', depth[idx], vmax=10.0)

            del inputs, outputs
//  [python]
def eval_save_all(self):
        """
        save prediction for all data on the list
        """
        self.set_eval()
        self.count = 0
        while True:
            try:
                inputs = self.val_iter.__next__()
            except StopIteration:
                break

            with torch.no_grad():
                inputs[("color_aug", 0, 0)] = inputs[("color_aug", 0, 0)].cuda()
                features = self.models["encoder"](inputs[("color_aug", 0, 0)]) #
                outputs = self.models["depth"](features)
                depth = output_to_depth(outputs[('out', 0)], self.opt.min_depth, self.opt.max_depth)
                outputs[("depth", 0, 0)] = depth
                sz = (640,640)
                store_path = f'results_all/'
                if not os.path.exists(store_path):
                    os.makedirs(store_path)
                    os.makedirs(store_path+'/image')
                    os.makedirs(store_path+'/depth')

                img = inputs[('color',0, 0)]
                img = t.interpolate(img, sz, mode='bilinear', align_corners=True)
                img = img.cpu().numpy().squeeze()

                if img.ndim == 3:
                    img = np.transpose(img, (1,2,0))
                else:
                    raise ValueError('Eval_save_all only supports batch_size = 1')

                if 'depth_gt' in inputs:
                    depth_gt = inputs['depth_gt']
                    depth_gt = t.interpolate(depth_gt, sz, mode='bilinear', align_corners=True)
                    depth_gt = depth_gt.cpu().numpy().squeeze()

                depth = outputs[('depth', 0, 0)] * self.approx_factor #approximate alignment for visualization
                depth = t.interpolate(depth, sz, mode='bilinear', align_corners=True)
                depth = depth.cpu().numpy().squeeze()

                self.count += 1
                imageio.imwrite(f'{store_path}/image/{self.count:04d}_img.png', img)
                write_turbo_depth_metric(f'{store_path}/depth/{self.count:04d}_depth.png', depth, vmax=10.0)
                if 'depth_gt' in inputs:
                    write_turbo_depth_metric(f'{store_path}/depth/{self.count:04}_depth_gt.png', depth_gt, vmax=10.0)

            del inputs, outputs
//  [python]
def eval_save_all(self):
        """
        save prediction for all data on the list
        """
        self.set_eval()
        self.count = 0
        while True:
            try:
                inputs = self.val_iter.__next__()
            except StopIteration:
                break

            with torch.no_grad():
                inputs[("color_aug", 0, 0)] = inputs[("color_aug", 0, 0)].cuda()
                features = self.models["encoder"](inputs[("color_aug", 0, 0)]) #
                outputs = self.models["depth"](features)
                depth = output_to_depth(outputs[('out', 0)], self.opt.min_depth, self.opt.max_depth)
                outputs[("depth", 0, 0)] = depth
                sz = (640,640)
                store_path = f'results_all/'
                if not os.path.exists(store_path):
                    os.makedirs(store_path)
                    os.makedirs(store_path+'/image')
                    os.makedirs(store_path+'/depth')

                img = inputs[('color',0, 0)]
                img = F.interpolate(img, sz, mode='bilinear', align_corners=True)
                img = img.cpu().numpy().squeeze()

                if img.ndim == 3:
                    img = np.transpose(img, (1,2,0))
                else:
                    raise ValueError('Eval_save_all only supports batch_size = 1')

                if 'depth_gt' in inputs:
                    depth_gt = inputs['depth_gt']
                    depth_gt = F.interpolate(depth_gt, sz, mode='bilinear', align_corners=True)
                    depth_gt = depth_gt.cpu().numpy().squeeze()

                depth = outputs[('depth', 0, 0)] * self.approx_factor #approximate alignment for visualization
                depth = F.interpolate(depth, sz, mode='bilinear', align_corners=True)
                depth = depth.cpu().numpy().squeeze()

                self.count += 1
                imageio.imwrite(f'{store_path}/image/{self.count:04d}_img.png', img)
                write_turbo_depth_metric(f'{store_path}/depth/{self.count:04d}_depth.png', depth, vmax=10.0)
                if 'depth_gt' in inputs:
                    write_turbo_depth_metric(f'{store_path}/depth/{self.count:04}_depth_gt.png', depth_gt, vmax=10.0)

            del inputs, outputs
//  [python]
def expand_f0(f0_frame, hop_length, method="interpolation"):
    """Expand f0 to output wave length.

    Args:
        f0_frame (Tensor): Input tensor (B, 1, frame_len).
        hop_length (Tensor): Hop length.
        method (str): Method to expand f0. Choose either 'interpolation' or 'repeat'.

    Returns:
        Tensor: Output tensor (B, 1, wav_len).

    """
    frame_length = f0_frame.size(2)
    signal_length = frame_length * hop_length
    if method == "interpolation":
        f0_sample = t.nn.functional.interpolate(
            f0_frame, size=signal_length, mode="linear", align_corners=False
        )
    elif method == "repeat":
        f0_sample = f0_frame.repeat_interleave(hop_length, dim=2)[:signal_length]
    else:
        raise ValueError("Invalid method. Choose either 'interpolation' or 'repeat'.")
    f0_sample = f0_sample.squeeze()[
        :signal_length
    ]  # Remove extra dimensions and trim to signal_length
    return f0_sample
//  [python]
def expand_f0(f0_frame, hop_length, method="interpolation"):
    """Expand f0 to output wave length.

    Args:
        f0_frame (Tensor): Input tensor (B, 1, frame_len).
        hop_length (Tensor): Hop length.
        method (str): Method to expand f0. Choose either 'interpolation' or 'repeat'.

    Returns:
        Tensor: Output tensor (B, 1, wav_len).

    """
    frame_length = f0_frame.size(2)
    signal_length = frame_length * hop_length
    if method == "interpolation":
        f0_sample = F.interpolate(
            f0_frame, size=signal_length, mode="linear", align_corners=False
        )
    elif method == "repeat":
        f0_sample = f0_frame.repeat_interleave(hop_length, dim=2)[:signal_length]
    else:
        raise ValueError("Invalid method. Choose either 'interpolation' or 'repeat'.")
    f0_sample = f0_sample.squeeze()[
        :signal_length
    ]  # Remove extra dimensions and trim to signal_length
    return f0_sample
//  [python]
def forward(self, x):
        sh = torch.tensor(x.shape)
        return F.interpolate(x, size=(int(sh[2]*self.scale), int(sh[3]*self.scale)), mode=self.mode, align_corners=self.align_corners)
//  [python]
def forward(self, inputs):
        B1, B2, B3 = inputs

        P3_x = self.P3_1(B3)
        P3_upsampled_x = F.upsample(P3_x, scale_factor=2)
        P3_x = self.P3_2(P3_x)

        P2_x = self.P2_1(B2)
        P2_x = P3_upsampled_x + P2_x
        P2_upsampled_x = F.upsample(P2_x, scale_factor=2)
        P2_x = self.P2_2(P2_x)

        P1_x = self.P1_1(B1)
        P1_x = P1_x + P2_upsampled_x
        P1_x = self.P1_2(P1_x)

        return [P1_x, P2_x, P3_x]
//  [python]
def forward(self, inputs):
        B1, B2, B3 = inputs

        P3_x = self.P3_1(B3)
        P3_upsampled_x = F.interpolate(P3_x, scale_factor=2)
        P3_x = self.P3_2(P3_x)

        P2_x = self.P2_1(B2)
        P2_x = P3_upsampled_x + P2_x
        P2_upsampled_x = F.interpolate(P2_x, scale_factor=2)
        P2_x = self.P2_2(P2_x)

        P1_x = self.P1_1(B1)
        P1_x = P1_x + P2_upsampled_x
        P1_x = self.P1_2(P1_x)

        return [P1_x, P2_x, P3_x]
//  [python]
def forward(self, prediction_dict, gt, ksi, scale):
        mask = (gt > 0).detach()

        output_size = prediction_dict["P", scale].size()
        depth_gt_scaled = t.nn.functional.interpolate(gt, size=(output_size[2], output_size[3]), mode='nearest')
        #p_gt, q_gt, r_gt, s_gt = split_depth2pqrs(depth_gt_scaled)
        p_gt, q_gt, r_gt, s_gt = depth2pqrs(depth_gt_scaled)

        p_prediction = prediction_dict["P", scale]
        q_prediction = prediction_dict["Q", scale]
        r_prediction = prediction_dict["R", scale]
        s_prediction = prediction_dict["S", scale]

        # Direction error - ranges from 0 (prediction is identical to GT) to pi (prediction is opposite from GT).
        # Problems with NaN: https://discuss.pytorch.org/t/nan-gradient-for-torch-cos-torch-acos/9617/2
        #direction_error = torch.acos(p_prediction * p_gt + q_prediction * q_gt + r_prediction * r_gt)
        inner_product = p_prediction * p_gt + q_prediction * q_gt + r_prediction * r_gt
        inner_product = torch.clamp(inner_product, min=-0.99, max=0.99)
        direction_error = torch.acos(inner_product)

        # Normal can be computed for eg. nx = -1 * normalized(p), as per literature
        # https://ieeexplore.ieee.org/document/7900061
        # nx_gt, ny_gt, nz_gt = -1*p_gt, -1*q_gt, -1*r_gt
        # nx_prediction, ny_prediction, nz_prediction = -1 * p_prediction, -1 * q_prediction, -1 * r_prediction
        # normal_dot_product = (nx_prediction * nx_gt + ny_prediction * ny_gt + nz_prediction * nz_gt)

        # direction_error = 1.0 - (p_prediction * p_gt + q_prediction * q_gt + r_prediction * r_gt)

        # Scale error - can grow unbounded with current implementation.
        scale_error = torch.maximum(torch.div(s_prediction, s_gt), torch.div(s_gt, s_prediction)) - 1

        direction_loss = torch.mean(direction_error[mask])
        scale_loss = torch.mean(scale_error[mask])

        return direction_loss + ksi * scale_loss
//  [python]
def forward(self, prediction_dict, gt, ksi, scale):
        mask = (gt > 0).detach()

        output_size = prediction_dict["P", scale].size()
        depth_gt_scaled = torch.nn.functional.interpolate(gt, size=(output_size[2], output_size[3]), mode='nearest')
        #p_gt, q_gt, r_gt, s_gt = split_depth2pqrs(depth_gt_scaled)
        p_gt, q_gt, r_gt, s_gt = depth2pqrs(depth_gt_scaled)

        p_prediction = prediction_dict["P", scale]
        q_prediction = prediction_dict["Q", scale]
        r_prediction = prediction_dict["R", scale]
        s_prediction = prediction_dict["S", scale]

        # Direction error - ranges from 0 (prediction is identical to GT) to pi (prediction is opposite from GT).
        # Problems with NaN: https://discuss.pytorch.org/t/nan-gradient-for-torch-cos-torch-acos/9617/2
        #direction_error = torch.acos(p_prediction * p_gt + q_prediction * q_gt + r_prediction * r_gt)
        inner_product = p_prediction * p_gt + q_prediction * q_gt + r_prediction * r_gt
        inner_product = torch.clamp(inner_product, min=-0.99, max=0.99)
        direction_error = torch.acos(inner_product)

        # Normal can be computed for eg. nx = -1 * normalized(p), as per literature
        # https://ieeexplore.ieee.org/document/7900061
        # nx_gt, ny_gt, nz_gt = -1*p_gt, -1*q_gt, -1*r_gt
        # nx_prediction, ny_prediction, nz_prediction = -1 * p_prediction, -1 * q_prediction, -1 * r_prediction
        # normal_dot_product = (nx_prediction * nx_gt + ny_prediction * ny_gt + nz_prediction * nz_gt)

        # direction_error = 1.0 - (p_prediction * p_gt + q_prediction * q_gt + r_prediction * r_gt)

        # Scale error - can grow unbounded with current implementation.
        scale_error = torch.maximum(torch.div(s_prediction, s_gt), torch.div(s_gt, s_prediction)) - 1

        direction_loss = torch.mean(direction_error[mask])
        scale_loss = torch.mean(scale_error[mask])

        return direction_loss + ksi * scale_loss
//  [python]
def forward(self, prediction_dict, gt, scale, pqrs2pqrs_loss=False):

        depth_prediction = prediction_dict["depth", scale]
        depth_gt_scaled = t.nn.functional.interpolate(gt,
                                                          size=(
                                                              depth_prediction.size()[2], depth_prediction.size()[3]),
                                                          mode='nearest')
        mask = (depth_gt_scaled > 0.1).detach()
        p_gt, q_gt, r_gt, s_gt = depth2pqrs(depth_gt_scaled)

        p_prediction = prediction_dict["P", scale]
        q_prediction = prediction_dict["Q", scale]
        r_prediction = prediction_dict["R", scale]
        s_prediction = prediction_dict["S", scale]

        if pqrs2pqrs_loss:

            po_prediction = prediction_dict["P_offset", scale]
            qo_prediction = prediction_dict["Q_offset", scale]
            ro_prediction = prediction_dict["R_offset", scale]
            so_prediction = prediction_dict["S_offset", scale]

            p_error = torch.abs(po_prediction[mask] - p_gt[mask])
            q_error = torch.abs(qo_prediction[mask] - q_gt[mask])
            r_error = torch.abs(ro_prediction[mask] - r_gt[mask])
            s_error = torch.abs(so_prediction[mask] - s_gt[mask])

        else:

            p_error = torch.abs(p_prediction[mask] - p_gt[mask])
            q_error = torch.abs(q_prediction[mask] - q_gt[mask])
            r_error = torch.abs(r_prediction[mask] - r_gt[mask])
            s_error = torch.abs(s_prediction[mask] - s_gt[mask])

        p_loss = torch.mean(p_error)
        q_loss = torch.mean(q_error)
        r_loss = torch.mean(r_error)
        s_loss = torch.mean(s_error)

        l1_loss = p_loss + q_loss + r_loss + s_loss

        return l1_loss
//  [python]
def forward(self, prediction_dict, gt, scale, pqrs2pqrs_loss=False):

        depth_prediction = prediction_dict["depth", scale]
        depth_gt_scaled = torch.nn.functional.interpolate(gt,
                                                          size=(
                                                              depth_prediction.size()[2], depth_prediction.size()[3]),
                                                          mode='nearest')
        mask = (depth_gt_scaled > 0.1).detach()
        p_gt, q_gt, r_gt, s_gt = depth2pqrs(depth_gt_scaled)

        p_prediction = prediction_dict["P", scale]
        q_prediction = prediction_dict["Q", scale]
        r_prediction = prediction_dict["R", scale]
        s_prediction = prediction_dict["S", scale]

        if pqrs2pqrs_loss:

            po_prediction = prediction_dict["P_offset", scale]
            qo_prediction = prediction_dict["Q_offset", scale]
            ro_prediction = prediction_dict["R_offset", scale]
            so_prediction = prediction_dict["S_offset", scale]

            p_error = torch.abs(po_prediction[mask] - p_gt[mask])
            q_error = torch.abs(qo_prediction[mask] - q_gt[mask])
            r_error = torch.abs(ro_prediction[mask] - r_gt[mask])
            s_error = torch.abs(so_prediction[mask] - s_gt[mask])

        else:

            p_error = torch.abs(p_prediction[mask] - p_gt[mask])
            q_error = torch.abs(q_prediction[mask] - q_gt[mask])
            r_error = torch.abs(r_prediction[mask] - r_gt[mask])
            s_error = torch.abs(s_prediction[mask] - s_gt[mask])

        p_loss = torch.mean(p_error)
        q_loss = torch.mean(q_error)
        r_loss = torch.mean(r_error)
        s_loss = torch.mean(s_error)

        l1_loss = p_loss + q_loss + r_loss + s_loss

        return l1_loss
//  [python]
def forward(self, feats):
        assert len(feats) == self.num_levels
        mask_feats = self.mask_feature_head(feats)
        feats = self.resize_feats(feats)
        mlvl_kernel_preds = []
        mlvl_cls_preds = []
        for i in range(self.num_levels):
            ins_kernel_feat = feats[i]
            # ins branch
            # concat coord
            coord_feat = generate_coordinate(ins_kernel_feat.size(),
                                             ins_kernel_feat.device)
            ins_kernel_feat = torch.cat([ins_kernel_feat, coord_feat], 1)

            # kernel branch
            kernel_feat = ins_kernel_feat
            kernel_feat = t.interpolate(
                kernel_feat,
                size=self.num_grids[i],
                mode='bilinear',
                align_corners=False)

            cate_feat = kernel_feat[:, :-2, :, :]

            kernel_feat = kernel_feat.contiguous()
            for i, kernel_conv in enumerate(self.kernel_convs):
                kernel_feat = kernel_conv(kernel_feat)
            kernel_pred = self.conv_kernel(kernel_feat)

            # cate branch
            cate_feat = cate_feat.contiguous()
            for i, cls_conv in enumerate(self.cls_convs):
                cate_feat = cls_conv(cate_feat)
            cate_pred = self.conv_cls(cate_feat)

            mlvl_kernel_preds.append(kernel_pred)
            mlvl_cls_preds.append(cate_pred)

        return mlvl_kernel_preds, mlvl_cls_preds, mask_feats
//  [python]
def forward(self, feats):
        assert len(feats) == self.num_levels
        mask_feats = self.mask_feature_head(feats)
        feats = self.resize_feats(feats)
        mlvl_kernel_preds = []
        mlvl_cls_preds = []
        for i in range(self.num_levels):
            ins_kernel_feat = feats[i]
            # ins branch
            # concat coord
            coord_feat = generate_coordinate(ins_kernel_feat.size(),
                                             ins_kernel_feat.device)
            ins_kernel_feat = torch.cat([ins_kernel_feat, coord_feat], 1)

            # kernel branch
            kernel_feat = ins_kernel_feat
            kernel_feat = F.interpolate(
                kernel_feat,
                size=self.num_grids[i],
                mode='bilinear',
                align_corners=False)

            cate_feat = kernel_feat[:, :-2, :, :]

            kernel_feat = kernel_feat.contiguous()
            for i, kernel_conv in enumerate(self.kernel_convs):
                kernel_feat = kernel_conv(kernel_feat)
            kernel_pred = self.conv_kernel(kernel_feat)

            # cate branch
            cate_feat = cate_feat.contiguous()
            for i, cls_conv in enumerate(self.cls_convs):
                cate_feat = cls_conv(cate_feat)
            cate_pred = self.conv_cls(cate_feat)

            mlvl_kernel_preds.append(kernel_pred)
            mlvl_cls_preds.append(cate_pred)

        return mlvl_kernel_preds, mlvl_cls_preds, mask_feats
//  [python]
def change_rms(data1, sr1, data2, sr2, rate):  # 1是输入音频，2是输出音频,rate是2的占比
    # print(data1.max(),data2.max())
    rms1 = librosa.feature.rms(
        y=data1, frame_length=sr1 // 2 * 2, hop_length=sr1 // 2
    )  # 每半秒一个点
    rms2 = librosa.feature.rms(y=data2, frame_length=sr2 // 2 * 2, hop_length=sr2 // 2)
    rms1 = torch.from_numpy(rms1)
    rms1 = t.upsample(
        rms1.unsqueeze(0), size=data2.shape[0], mode="linear"
    ).squeeze()
    rms2 = torch.from_numpy(rms2)
    rms2 = t.upsample(
        rms2.unsqueeze(0), size=data2.shape[0], mode="linear"
    ).squeeze()
    rms2 = torch.max(rms2, torch.zeros_like(rms2) + 1e-6)
    data2 *= (
        torch.pow(rms1, torch.tensor(1 - rate))
        * torch.pow(rms2, torch.tensor(rate - 1))
    ).numpy()
    return data2
//  [python]
def change_rms(data1, sr1, data2, sr2, rate):  # 1是输入音频，2是输出音频,rate是2的占比
    # print(data1.max(),data2.max())
    rms1 = librosa.feature.rms(
        y=data1, frame_length=sr1 // 2 * 2, hop_length=sr1 // 2
    )  # 每半秒一个点
    rms2 = librosa.feature.rms(y=data2, frame_length=sr2 // 2 * 2, hop_length=sr2 // 2)
    rms1 = torch.from_numpy(rms1)
    rms1 = F.interpolate(
        rms1.unsqueeze(0), size=data2.shape[0], mode="linear"
    ).squeeze()
    rms2 = torch.from_numpy(rms2)
    rms2 = F.interpolate(
        rms2.unsqueeze(0), size=data2.shape[0], mode="linear"
    ).squeeze()
    rms2 = torch.max(rms2, torch.zeros_like(rms2) + 1e-6)
    data2 *= (
        torch.pow(rms1, torch.tensor(1 - rate))
        * torch.pow(rms2, torch.tensor(rate - 1))
    ).numpy()
    return data2