jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def visualize_results_vps(self, batched_inputs, pan_seg_dict, out_size, sub_task):
        '''
        pan_seg_dict = {
            "image_size": out_size,
            "pred_masks": panoptic_seg.cpu(),  # V x H x W
            "segments_infos": segments_infos,
            "task": "vps",
            }
        '''
        file_names = batched_inputs[0]['file_names'] 
        dataset_name = file_names[0].split('/')[1]
        video_id = file_names[0].split('/')[-2]
        video_len = len(file_names)

        save_dir = os.path.join(self.output_dir, 'inference/vps_'+dataset_name, video_id)
        os.makedirs(save_dir, exist_ok=True)

        pan_segs = pan_seg_dict["pred_masks"]
        segm_infos = pan_seg_dict["segments_infos"]
        for segm_info in segm_infos:
            if segm_info['isthing']:
                segm_info['category_id'] = self.metadata.thing_dataset_id_to_contiguous_id[segm_info['category_id']]
            else:
                segm_info['category_id'] = self.metadata.stuff_dataset_id_to_contiguous_id[segm_info['category_id']]
        
        assert video_len == len(pan_segs), f'Mismatch length between predicted and GT masks: {len(pan_segs)} and {video_len}'
        for t, pan_seg in enumerate(pan_segs):
            file_name = file_names[t].split('/')[-1]
            img = batched_inputs[0]["image"][t]
            img = t.interpolate(
                img[None].float(),
                out_size,
                mode="bilinear",
                align_corners=False
            ).squeeze(0).long()
            img = img.permute(1, 2, 0).cpu().to(torch.uint8)
            visualizer = VisualizerFrame(
                img, 
                metadata=self.metadata, 
                class_names=None
            )

            save_path = '/'.join([save_dir, file_name.replace('.png', '.ipg')])
            VisImage = visualizer.draw_panoptic_seg(pan_seg, segm_infos, alpha=0.8)
            VisImage.save(save_path)
        
        # save all frames to a .avi video
        out = cv2.VideoWriter(
            '/'.join([save_dir, video_id + '.avi']),
            cv2.VideoWriter_fourcc(*'DIVX'),
            fps=10,
            frameSize=(out_size[1], out_size[0])
        )

        file_names = glob.glob('/'.join([save_dir, "*.jpg"]))
        file_names = sorted(
            file_names, 
            key=lambda f: int(f.split("/")[-1].split('_')[-1].replace('.jpg', ''))
        )
        for file_name in file_names:
            out.write(cv2.imread(file_name))
        out.release()
        print(f"save all frames with pan. seg. into {video_id}.avi")
//  [python]
def visualize_results_vps(self, batched_inputs, pan_seg_dict, out_size, sub_task):
        '''
        pan_seg_dict = {
            "image_size": out_size,
            "pred_masks": panoptic_seg.cpu(),  # V x H x W
            "segments_infos": segments_infos,
            "task": "vps",
            }
        '''
        file_names = batched_inputs[0]['file_names'] 
        dataset_name = file_names[0].split('/')[1]
        video_id = file_names[0].split('/')[-2]
        video_len = len(file_names)

        save_dir = os.path.join(self.output_dir, 'inference/vps_'+dataset_name, video_id)
        os.makedirs(save_dir, exist_ok=True)

        pan_segs = pan_seg_dict["pred_masks"]
        segm_infos = pan_seg_dict["segments_infos"]
        for segm_info in segm_infos:
            if segm_info['isthing']:
                segm_info['category_id'] = self.metadata.thing_dataset_id_to_contiguous_id[segm_info['category_id']]
            else:
                segm_info['category_id'] = self.metadata.stuff_dataset_id_to_contiguous_id[segm_info['category_id']]
        
        assert video_len == len(pan_segs), f'Mismatch length between predicted and GT masks: {len(pan_segs)} and {video_len}'
        for t, pan_seg in enumerate(pan_segs):
            file_name = file_names[t].split('/')[-1]
            img = batched_inputs[0]["image"][t]
            img = F.interpolate(
                img[None].float(),
                out_size,
                mode="bilinear",
                align_corners=False
            ).squeeze(0).long()
            img = img.permute(1, 2, 0).cpu().to(torch.uint8)
            visualizer = VisualizerFrame(
                img, 
                metadata=self.metadata, 
                class_names=None
            )

            save_path = '/'.join([save_dir, file_name.replace('.png', '.ipg')])
            VisImage = visualizer.draw_panoptic_seg(pan_seg, segm_infos, alpha=0.8)
            VisImage.save(save_path)
        
        # save all frames to a .avi video
        out = cv2.VideoWriter(
            '/'.join([save_dir, video_id + '.avi']),
            cv2.VideoWriter_fourcc(*'DIVX'),
            fps=10,
            frameSize=(out_size[1], out_size[0])
        )

        file_names = glob.glob('/'.join([save_dir, "*.jpg"]))
        file_names = sorted(
            file_names, 
            key=lambda f: int(f.split("/")[-1].split('_')[-1].replace('.jpg', ''))
        )
        for file_name in file_names:
            out.write(cv2.imread(file_name))
        out.release()
        print(f"save all frames with pan. seg. into {video_id}.avi")
//  [python]
def forward(self, inputs,outs=[]):
        """
        :param: inputs+ upsampled seg maps at 4 scales
        :return: output map at 4 scales
        """
        scale_num=len(outs)
        outs_ref=[]
        for i in range(scale_num-1):
            cur_inputs=t.interpolate(inputs,scale_factor=1/(2**(3-i)), mode='bilinear', align_corners=True)
            if i==0:
                unet_inputs=torch.cat([cur_inputs,outs[i],outs[i]],dim=1)
                cur_outs=self.ref_block(unet_inputs)
                outs_ref.append(cur_outs)
            else:
                cur_outs=t.interpolate(cur_outs,scale_factor=2, mode='bilinear', align_corners=True)
                unet_inputs = torch.cat([cur_inputs, outs[i], cur_outs], dim=1)
                cur_outs = self.ref_block(unet_inputs)
                outs_ref.append(cur_outs)
        cur_outs = t.interpolate(cur_outs, scale_factor=2, mode='bilinear', align_corners=True)
        unet_inputs = torch.cat([inputs, outs[-1], cur_outs], dim=1)
        cur_outs = self.ref_block(unet_inputs)
        outs_ref.append(cur_outs)


        return  outs_ref
//  [python]
def forward(self, inputs,outs=[]):
        """
        :param: inputs+ upsampled seg maps at 4 scales
        :return: output map at 4 scales
        """
        scale_num=len(outs)
        outs_ref=[]
        for i in range(scale_num-1):
            cur_inputs=F.interpolate(inputs,scale_factor=1/(2**(3-i)), mode='bilinear', align_corners=True)
            if i==0:
                unet_inputs=torch.cat([cur_inputs,outs[i],outs[i]],dim=1)
                cur_outs=self.ref_block(unet_inputs)
                outs_ref.append(cur_outs)
            else:
                cur_outs=F.interpolate(cur_outs,scale_factor=2, mode='bilinear', align_corners=True)
                unet_inputs = torch.cat([cur_inputs, outs[i], cur_outs], dim=1)
                cur_outs = self.ref_block(unet_inputs)
                outs_ref.append(cur_outs)
        cur_outs = F.interpolate(cur_outs, scale_factor=2, mode='bilinear', align_corners=True)
        unet_inputs = torch.cat([inputs, outs[-1], cur_outs], dim=1)
        cur_outs = self.ref_block(unet_inputs)
        outs_ref.append(cur_outs)


        return  outs_ref
//  [python]
def forward(self, inputs):
        inputs0 = t.interpolate(inputs, scale_factor=0.5, mode='bilinear', align_corners=True)
        out_logits0 = self.segModel(inputs0)
        out_logits1 = self.segModel(inputs)

        out_logits00=t.interpolate(out_logits0, scale_factor=2, mode='bilinear', align_corners=True)
        out_logits01=torch.cat([out_logits00,out_logits1],dim=1)

        att_map01 = self.scale_head(out_logits01)

        return out_logits00,out_logits1,att_map01
//  [python]
def forward(self, inputs):
        inputs0 = F.interpolate(inputs, scale_factor=0.5, mode='bilinear', align_corners=True)
        out_logits0 = self.segModel(inputs0)
        out_logits1 = self.segModel(inputs)

        out_logits00=F.interpolate(out_logits0, scale_factor=2, mode='bilinear', align_corners=True)
        out_logits01=torch.cat([out_logits00,out_logits1],dim=1)

        att_map01 = self.scale_head(out_logits01)

        return out_logits00,out_logits1,att_map01
//  [python]
def forward(self, x):
        gap = self.gap(x)
        out = x + t.upsample_bilinear(gap, size=x.shape[2:])
        out = self.conv(out)
        return out
//  [python]
def forward(self, x):
        gap = self.gap(x)
        out = x + F.interpolate(gap, x.shape[2:], mode='bilinear', align_corners=True)
        out = self.conv(out)
        return out
//  [python]
def forward(self, x, size=None):
        out = self.D(x)
        src_out = self.cls1(out)
        tgt_out = self.cls2(out)
        out = torch.cat((src_out, tgt_out), dim=1)
        if size is not None:
            out = t.interpolate(out, size=size, mode='bilinear', align_corners=True)
        return out
//  [python]
def forward(self, x, size=None):
        out = self.D(x)
        src_out = self.cls1(out)
        tgt_out = self.cls2(out)
        out = torch.cat((src_out, tgt_out), dim=1)
        if size is not None:
            out = F.interpolate(out, size=size, mode='bilinear', align_corners=True)
        return out
//  [python]
def forward(self, x):
        h, w = x.shape[2:]

        if self.training or self.model_name != "icnet":  # general settings or pspnet
            k_sizes = []
            strides = []
            for pool_size in self.pool_sizes:
                k_sizes.append((int(h / pool_size), int(w / pool_size)))
                strides.append((int(h / pool_size), int(w / pool_size)))
        else:  # eval mode and icnet: pre-trained for 1025 x 2049
            k_sizes = [(8, 15), (13, 25), (17, 33), (33, 65)]
            strides = [(5, 10), (10, 20), (16, 32), (33, 65)]

        if self.fusion_mode == "cat":  # pspnet: concat (including x)
            output_slices = [x]

            for i, (module, pool_size) in enumerate(zip(self.path_module_list, self.pool_sizes)):
                out = F.avg_pool2d(x, k_sizes[i], stride=strides[i], padding=0)
                # out = F.adaptive_avg_pool2d(x, output_size=(pool_size, pool_size))
                if self.model_name != "icnet":
                    out = module(out)
                out = t.upsample(out, size=(h, w), mode="bilinear")
                output_slices.append(out)

            return torch.cat(output_slices, dim=1)
        else:  # icnet: element-wise sum (including x)
            pp_sum = x

            for i, (module, pool_size) in enumerate(zip(self.path_module_list, self.pool_sizes)):
                out = F.avg_pool2d(x, k_sizes[i], stride=strides[i], padding=0)
                # out = F.adaptive_avg_pool2d(x, output_size=(pool_size, pool_size))
                if self.model_name != "icnet":
                    out = module(out)
                out = t.upsample(out, size=(h, w), mode="bilinear")
                pp_sum = pp_sum + out

            return pp_sum
//  [python]
def forward(self, x):
        h, w = x.shape[2:]

        if self.training or self.model_name != "icnet":  # general settings or pspnet
            k_sizes = []
            strides = []
            for pool_size in self.pool_sizes:
                k_sizes.append((int(h / pool_size), int(w / pool_size)))
                strides.append((int(h / pool_size), int(w / pool_size)))
        else:  # eval mode and icnet: pre-trained for 1025 x 2049
            k_sizes = [(8, 15), (13, 25), (17, 33), (33, 65)]
            strides = [(5, 10), (10, 20), (16, 32), (33, 65)]

        if self.fusion_mode == "cat":  # pspnet: concat (including x)
            output_slices = [x]

            for i, (module, pool_size) in enumerate(zip(self.path_module_list, self.pool_sizes)):
                out = F.avg_pool2d(x, k_sizes[i], stride=strides[i], padding=0)
                # out = F.adaptive_avg_pool2d(x, output_size=(pool_size, pool_size))
                if self.model_name != "icnet":
                    out = module(out)
                out = F.interpolate(out, size=(h, w), mode="bilinear", align_corners=True)
                output_slices.append(out)

            return torch.cat(output_slices, dim=1)
        else:  # icnet: element-wise sum (including x)
            pp_sum = x

            for i, (module, pool_size) in enumerate(zip(self.path_module_list, self.pool_sizes)):
                out = F.avg_pool2d(x, k_sizes[i], stride=strides[i], padding=0)
                # out = F.adaptive_avg_pool2d(x, output_size=(pool_size, pool_size))
                if self.model_name != "icnet":
                    out = module(out)
                out = F.interpolate(out, size=(h, w), mode="bilinear", align_corners=True)
                pp_sum = pp_sum + out

            return pp_sum
//  [python]
def forward(self, x):
        outs = super().forward(x)
        if len(outs) > 1:
            resize_outs = []
            if self.use_adp:
                for i in range(len(outs)):
                    feature = self.adp[i](outs[i])
                    resize_outs.append(feature)
            else:
                target_size = self.target_size
                for i in range(len(outs)):
                    feature = outs[i]
                    if feature.shape[2:] != target_size:
                        feature = t.interpolate(feature, target_size,  mode='bilinear', align_corners=True)
                    resize_outs.append(feature)
            out = torch.cat(resize_outs, dim=1)
            out = self.reduc_conv(out)
                

        else:
            out = outs[0]
        return [out]
//  [python]
def forward(self, x):
        outs = super().forward(x)
        if len(outs) > 1:
            resize_outs = []
            if self.use_adp:
                for i in range(len(outs)):
                    feature = self.adp[i](outs[i])
                    resize_outs.append(feature)
            else:
                target_size = self.target_size
                for i in range(len(outs)):
                    feature = outs[i]
                    if feature.shape[2:] != target_size:
                        feature = F.interpolate(feature, target_size,  mode='bilinear', align_corners=True)
                    resize_outs.append(feature)
            out = torch.cat(resize_outs, dim=1)
            out = self.reduc_conv(out)
                

        else:
            out = outs[0]
        return [out]
//  [python]
def forward(self, x):
        outputs = t.interpolate(
            x,
            size=x.shape[-1] * self.scale_factor \
                if not self.downsample else x.shape[-1] // self.scale_factor,
            mode=self.mode,
            align_corners=self.align_corners
        )
        return outputs
//  [python]
def forward(self, x):
        outputs = torch.nn.functional.interpolate(
            x,
            size=x.shape[-1] * self.scale_factor \
                if not self.downsample else x.shape[-1] // self.scale_factor,
            mode=self.mode,
            align_corners=self.align_corners,
            recompute_scale_factor=False
        )
        return outputs
//  [python]
def test_interpolate_no_shape(self):
        class MyModel(torch.jit.ScriptModule):
            @torch.jit.script_method
            def forward(self, x, y):
                x = torch.add(x, x)
                out1 = torch.nn.functional.interpolate(x, mode="bilinear", size=(16, 16), align_corners=False)
                out2 = torch.nn.functional.interpolate(x, mode="nearest", size=(int(y.size(0)), int(y.size(1))))
                return out1, out2

        x = torch.randn(1, 2, 4, 4, requires_grad=True)
        y = torch.randn(16, 16, requires_grad=True)
        self.run_test(MyModel(), (x, y))
//  [python]
def test_interpolate_function_substitution(self):
        class ScriptModel(torch.jit.ScriptModule):
            @torch.jit.script_method
            def forward(self, x):
                return t.interpolate(x, mode="nearest", scale_factor=2.)

        class ScriptModule(torch.jit.ScriptModule):
            def __init__(self):
                super(ScriptModule, self).__init__()
                self.submodule = ScriptModel()

            @torch.jit.script_method
            def forward(self, input):
                return self.submodule(input)

        x = torch.randn(1, 2, 4, 4, 6)
        self.run_test(ScriptModule(), (x,))

        @torch.jit.script
        def script_method(x):
            return t.interpolate(x, mode="nearest", scale_factor=2.)

        class TracingModule(torch.nn.Module):
            def forward(self, x):
                return script_method(x)

        self.run_test(TracingModule(), (x,))
//  [python]
def test_interpolate_function_substitution(self):
        class ScriptModel(torch.jit.ScriptModule):
            @torch.jit.script_method
            def forward(self, x):
                return torch.nn.functional.interpolate(x, mode="nearest", scale_factor=2.)

        class ScriptModule(torch.jit.ScriptModule):
            def __init__(self):
                super(ScriptModule, self).__init__()
                self.submodule = ScriptModel()

            @torch.jit.script_method
            def forward(self, input):
                return self.submodule(input)

        x = torch.randn(1, 2, 4, 4, 6)
        self.run_test(ScriptModule(), (x,))

        @torch.jit.script
        def script_method(x):
            return torch.nn.functional.interpolate(x, mode="nearest", scale_factor=2.)

        class TracingModule(torch.nn.Module):
            def forward(self, x):
                return script_method(x)

        self.run_test(TracingModule(), (x,))
//  [python]
def _interpolate(self, x, mode, use_size, is_upsample, align_corners=False):
        class MyModel(torch.nn.Module):
            def forward(self, x):
                scale = 2.3 if is_upsample else 0.5
                if len(x.size()) == 3:
                    scale_array = 2.3
                if len(x.size()) == 4:
                    scale_array = [2.3, 5.1]
                if len(x.size()) == 5:
                    scale_array = [3.3, 2.3, 5.1]
                if use_size:
                    size_array = [int(float(v) * scale) for v in x.size()[2:]]
                    if align_corners:
                        return t.interpolate(x, mode=mode, size=size_array[0], align_corners=True), \
                            t.interpolate(x, mode=mode, size=size_array, align_corners=True)
                    return t.interpolate(x, mode=mode, size=size_array[0]), \
                        t.interpolate(x, mode=mode, size=size_array)
                if align_corners:
                    return t.interpolate(x, mode=mode, scale_factor=scale,
                                                           align_corners=True), \
                        t.interpolate(x, mode=mode, scale_factor=scale_array,
                                                        align_corners=True)
                return t.interpolate(x, mode=mode,
                                                       scale_factor=scale), \
                    t.interpolate(x, mode=mode,
                                                    scale_factor=scale_array)

        self.run_test(MyModel(), x)
//  [python]
def _interpolate(self, x, mode, use_size, is_upsample, align_corners=False):
        class MyModel(torch.nn.Module):
            def forward(self, x):
                scale = 2.3 if is_upsample else 0.5
                if len(x.size()) == 3:
                    scale_array = 2.3
                if len(x.size()) == 4:
                    scale_array = [2.3, 5.1]
                if len(x.size()) == 5:
                    scale_array = [3.3, 2.3, 5.1]
                if use_size:
                    size_array = [int(float(v) * scale) for v in x.size()[2:]]
                    if align_corners:
                        return torch.nn.functional.interpolate(x, mode=mode, size=size_array[0], align_corners=True), \
                            torch.nn.functional.interpolate(x, mode=mode, size=size_array, align_corners=True)
                    return torch.nn.functional.interpolate(x, mode=mode, size=size_array[0]), \
                        torch.nn.functional.interpolate(x, mode=mode, size=size_array)
                if align_corners:
                    return torch.nn.functional.interpolate(x, mode=mode, scale_factor=scale,
                                                           align_corners=True, recompute_scale_factor=False), \
                        torch.nn.functional.interpolate(x, mode=mode, scale_factor=scale_array,
                                                        align_corners=True, recompute_scale_factor=False)
                return torch.nn.functional.interpolate(x, mode=mode,
                                                       scale_factor=scale, recompute_scale_factor=False), \
                    torch.nn.functional.interpolate(x, mode=mode,
                                                    scale_factor=scale_array, recompute_scale_factor=False)

        self.run_test(MyModel(), x)
//  [python]
def forward(self, x):
        #print("check ghostv2 input size:",x.size())
        if self.mode in ['original']:
            x1 = self.primary_conv(x)
            x2 = self.cheap_operation(x1)
            out = torch.cat([x1, x2], dim=1)
            return out[:, :self.oup, :, :]
        elif self.mode in ['attn']:
            res = self.short_conv(F.avg_pool2d(x, kernel_size=2, stride=2))
            x1 = self.primary_conv(x)
            x2 = self.cheap_operation(x1)
            out = torch.cat([x1, x2], dim=1)
            out=out[:, :self.oup, :, :] * F.upsample(self.gate_fn(res), size=(out.shape[-2], out.shape[-1]),
                                                           mode='nearest')
            return out
//  [python]
def forward(self, x):
        #print("check ghostv2 input size:",x.size())
        if self.mode in ['original']:
            x1 = self.primary_conv(x)
            x2 = self.cheap_operation(x1)
            out = torch.cat([x1, x2], dim=1)
            return out[:, :self.oup, :, :]
        elif self.mode in ['attn']:
            res = self.short_conv(F.avg_pool2d(x, kernel_size=2, stride=2))
            x1 = self.primary_conv(x)
            x2 = self.cheap_operation(x1)
            out = torch.cat([x1, x2], dim=1)
            out=out[:, :self.oup, :, :] * F.interpolate(self.gate_fn(res), size=(out.shape[-2], out.shape[-1]),
                                                           mode='nearest')
            return out
//  [python]
def forward(self, xl, xh):
        xl1 = self.lbody(xl)
        xl1 = F.upsample(
            xl1, size=xh.size()[2:], mode='bilinear')
        xh1 = self.hbody(xh)
        x = torch.cat((xl1, xh1), dim=1)
        x_out = self.body(x)
        return x_out
//  [python]
def forward(self, xl, xh):
        xl1 = self.lbody(xl)
        xl1 = F.interpolate(
            xl1, size=xh.size()[2:], mode='bilinear', align_corners=True)
        xh1 = self.hbody(xh)
        x = torch.cat((xl1, xh1), dim=1)
        x_out = self.body(x)
        return x_out
//  [python]
def forward(self, xl, xh):
        xl1 = self.lbody(xl)
        xh1 = self.hbody(xh)
        xh1 = F.upsample(
            xh1, size=xl.size()[2:], mode='bilinear')
        x = torch.cat((xl1, xh1), dim=1)
        x_out = self.bodye(x)
        return x_out
//  [python]
def forward(self, xl, xh):
        xl1 = self.lbody(xl)
        xh1 = self.hbody(xh)
        xh1 = F.interpolate(
            xh1, size=xl.size()[2:], mode='bilinear', align_corners=True)
        x = torch.cat((xl1, xh1), dim=1)
        x_out = self.bodye(x)
        return x_out
//  [python]
def repeat_expand(content: Union[torch.Tensor, np.ndarray], target_len: int, mode: str = "nearest"):
    """Repeat content to target length.
    This is a wrapper of torch.nn.functional.interpolate.

    Args:
        content (torch.Tensor): tensor
        target_len (int): target length
        mode (str, optional): interpolation mode. Defaults to "nearest".

    Returns:
        torch.Tensor: tensor
    """

    ndim = content.ndim

    if content.ndim == 1:
        content = content[None, None]
    elif content.ndim == 2:
        content = content[None]

    assert content.ndim == 3

    is_np = isinstance(content, np.ndarray)
    if is_np:
        content = torch.from_numpy(content)

    results = t.interpolate(content, size=target_len, mode=mode)

    if is_np:
        results = results.numpy()

    if ndim == 1:
        return results[0, 0]
    elif ndim == 2:
        return results[0]
//  [python]
def repeat_expand(content: Union[torch.Tensor, np.ndarray], target_len: int, mode: str = "nearest"):
    """Repeat content to target length.
    This is a wrapper of torch.nn.functional.interpolate.

    Args:
        content (torch.Tensor): tensor
        target_len (int): target length
        mode (str, optional): interpolation mode. Defaults to "nearest".

    Returns:
        torch.Tensor: tensor
    """

    ndim = content.ndim

    if content.ndim == 1:
        content = content[None, None]
    elif content.ndim == 2:
        content = content[None]

    assert content.ndim == 3

    is_np = isinstance(content, np.ndarray)
    if is_np:
        content = torch.from_numpy(content)

    results = torch.nn.functional.interpolate(content, size=target_len, mode=mode)

    if is_np:
        results = results.numpy()

    if ndim == 1:
        return results[0, 0]
    elif ndim == 2:
        return results[0]
//  [python]
def interpolate(
    input, size=None, scale_factor=None, mode="nearest", align_corners=None
):
    if input.numel() > 0:
        return t.nn.functional.interpolate(
            input, size, scale_factor, mode, align_corners
        )

    def _check_size_scale_factor(dim):
        if size is None and scale_factor is None:
            raise ValueError("either size or scale_factor should be defined")
        if size is not None and scale_factor is not None:
            raise ValueError("only one of size or scale_factor should be defined")
        if (
            scale_factor is not None
            and isinstance(scale_factor, tuple)
            and len(scale_factor) != dim
        ):
            raise ValueError(
                "scale_factor shape must match input shape. "
                "Input is {}D, scale_factor size is {}".format(dim, len(scale_factor))
            )

    def _output_size(dim):
        _check_size_scale_factor(dim)
        if size is not None:
            return size
        scale_factors = _ntuple(dim)(scale_factor)
        # math.floor might return float in py2.7
        return [
            int(math.floor(input.size(i + 2) * scale_factors[i])) for i in range(dim)
        ]

    output_shape = tuple(_output_size(2))
    output_shape = input.shape[:-2] + output_shape
    return _NewEmptyTensorOp.apply(input, output_shape)
//  [python]
def interpolate(
    input, size=None, scale_factor=None, mode="nearest", align_corners=None
):
    if input.numel() > 0:
        return torch.nn.functional.interpolate(
            input, size, scale_factor, mode, align_corners
        )

    def _check_size_scale_factor(dim):
        if size is None and scale_factor is None:
            raise ValueError("either size or scale_factor should be defined")
        if size is not None and scale_factor is not None:
            raise ValueError("only one of size or scale_factor should be defined")
        if (
            scale_factor is not None
            and isinstance(scale_factor, tuple)
            and len(scale_factor) != dim
        ):
            raise ValueError(
                "scale_factor shape must match input shape. "
                "Input is {}D, scale_factor size is {}".format(dim, len(scale_factor))
            )

    def _output_size(dim):
        _check_size_scale_factor(dim)
        if size is not None:
            return size
        scale_factors = _ntuple(dim)(scale_factor)
        # math.floor might return float in py2.7
        return [
            int(math.floor(input.size(i + 2) * scale_factors[i])) for i in range(dim)
        ]

    output_shape = tuple(_output_size(2))
    output_shape = input.shape[:-2] + output_shape
    return _NewEmptyTensorOp.apply(input, output_shape)
//  [python]
def forward(self, x):
        conv0 = self.conv0(x)
        conv1 = self.conv1(conv0)
        conv2 = self.conv2(conv1)

        intra_feat = conv2
        outputs = {}
        out = self.out1(intra_feat)
        outputs['stage1'] = out
        if self.arch_mode == 'unet':
            if self.num_stage == 3:
                intra_feat = self.deconv1(conv1, intra_feat)
                out = self.out2(intra_feat)
                outputs['stage2'] = out

                intra_feat = self.deconv2(conv0, intra_feat)
                out = self.out3(intra_feat)
                outputs['stage3'] = out

            elif self.num_stage == 2:
                intra_feat = self.deconv1(conv1, intra_feat)
                out = self.out2(intra_feat)
                outputs['stage2'] = out

        elif self.arch_mode == 'fpn':
            if self.num_stage == 3:
                intra_feat = t.interpolate(
                    intra_feat, scale_factor=2,
                    mode='nearest') + self.inner1(conv1)
                out = self.out2(intra_feat)
                outputs['stage2'] = out

                intra_feat = t.interpolate(
                    intra_feat, scale_factor=2,
                    mode='nearest') + self.inner2(conv0)
                out = self.out3(intra_feat)
                outputs['stage3'] = out

            elif self.num_stage == 2:
                intra_feat = t.interpolate(
                    intra_feat, scale_factor=2,
                    mode='nearest') + self.inner1(conv1)
                out = self.out2(intra_feat)
                outputs['stage2'] = out

        return outputs
//  [python]
def forward(self, x):
        conv0 = self.conv0(x)
        conv1 = self.conv1(conv0)
        conv2 = self.conv2(conv1)

        intra_feat = conv2
        outputs = {}
        out = self.out1(intra_feat)
        outputs['stage1'] = out
        if self.arch_mode == 'unet':
            if self.num_stage == 3:
                intra_feat = self.deconv1(conv1, intra_feat)
                out = self.out2(intra_feat)
                outputs['stage2'] = out

                intra_feat = self.deconv2(conv0, intra_feat)
                out = self.out3(intra_feat)
                outputs['stage3'] = out

            elif self.num_stage == 2:
                intra_feat = self.deconv1(conv1, intra_feat)
                out = self.out2(intra_feat)
                outputs['stage2'] = out

        elif self.arch_mode == 'fpn':
            if self.num_stage == 3:
                intra_feat = F.interpolate(
                    intra_feat, scale_factor=2,
                    mode='nearest') + self.inner1(conv1)
                out = self.out2(intra_feat)
                outputs['stage2'] = out

                intra_feat = F.interpolate(
                    intra_feat, scale_factor=2,
                    mode='nearest') + self.inner2(conv0)
                out = self.out3(intra_feat)
                outputs['stage3'] = out

            elif self.num_stage == 2:
                intra_feat = F.interpolate(
                    intra_feat, scale_factor=2,
                    mode='nearest') + self.inner1(conv1)
                out = self.out2(intra_feat)
                outputs['stage2'] = out

        return outputs
//  [python]
def forward(self, waveforms):
        """
        Arguments
        ---------
        waveforms : torch.Tensor
            Shape should be `[batch, time]` or `[batch, time, channels]`.

        Returns
        -------
        Tensor of shape `[batch, time]` or `[batch, time, channels]`.
        """

        if self.reverb_sample_rate != self.clean_sample_rate:
            self.resampler = Resample(
                self.reverb_sample_rate, self.clean_sample_rate
            )

        # Add channels dimension if necessary
        channel_added = False
        if len(waveforms.shape) == 2:
            waveforms = waveforms.unsqueeze(-1)
            channel_added = True

        # Load and prepare RIR
        rir_waveform = self._load_rir(waveforms)

        # Resample to correct rate
        if hasattr(self, "resampler"):
            rir_waveform = self.resampler(rir_waveform)

        # Compress or dilate RIR
        if self.rir_scale_factor != 1:
            rir_waveform = t.nn.functional.interpolate(
                rir_waveform.transpose(1, -1),
                scale_factor=self.rir_scale_factor,
                mode="linear",
                align_corners=False,
            )
            rir_waveform = rir_waveform.transpose(1, -1)

        rev_waveform = reverberate(waveforms, rir_waveform, rescale_amp="avg")

        # Remove channels dimension if added
        if channel_added:
            return rev_waveform.squeeze(-1)

        return rev_waveform
//  [python]
def forward(self, waveforms):
        """
        Arguments
        ---------
        waveforms : torch.Tensor
            Shape should be `[batch, time]` or `[batch, time, channels]`.

        Returns
        -------
        Tensor of shape `[batch, time]` or `[batch, time, channels]`.
        """

        if self.reverb_sample_rate != self.clean_sample_rate:
            self.resampler = Resample(
                self.reverb_sample_rate, self.clean_sample_rate
            )

        # Add channels dimension if necessary
        channel_added = False
        if len(waveforms.shape) == 2:
            waveforms = waveforms.unsqueeze(-1)
            channel_added = True

        # Load and prepare RIR
        rir_waveform = self._load_rir(waveforms)

        # Resample to correct rate
        if hasattr(self, "resampler"):
            rir_waveform = self.resampler(rir_waveform)

        # Compress or dilate RIR
        if self.rir_scale_factor != 1:
            rir_waveform = F.interpolate(
                rir_waveform.transpose(1, -1),
                scale_factor=self.rir_scale_factor,
                mode="linear",
                align_corners=False,
            )
            rir_waveform = rir_waveform.transpose(1, -1)

        rev_waveform = reverberate(waveforms, rir_waveform, rescale_amp="avg")

        # Remove channels dimension if added
        if channel_added:
            return rev_waveform.squeeze(-1)

        return rev_waveform
//  [python]
def forward(self, x, edgeAtten):
        if edgeAtten.size() != x.size():
            edgeAtten = F.upsample(
                edgeAtten, x.size()[2:], mode='bilinear')
        xx = torch.chunk(x, self.ne_modules, dim=1)
        efeas = []
        for i in range(self.ne_modules):
            xei = self.emlayes[i](xx[i], edgeAtten)
            efeas.append(xei)
        efeas = torch.cat(efeas, dim=1)
        x_out = self.body(efeas)
        return x_out
//  [python]
def forward(self, x, edgeAtten):
        if edgeAtten.size() != x.size():
            edgeAtten = F.interpolate(
                edgeAtten, x.size()[2:], mode='bilinear', align_corners=False)
        xx = torch.chunk(x, self.ne_modules, dim=1)
        efeas = []
        for i in range(self.ne_modules):
            xei = self.emlayes[i](xx[i], edgeAtten)
            efeas.append(xei)
        efeas = torch.cat(efeas, dim=1)
        x_out = self.body(efeas)
        return x_out
//  [python]
def forward(self, img):

        img = self.sub_mean(img)

        f1=self.conv3x3(img)
        f2=self.DG1(f1)
        f3=self.DG2(self.conv1x1_1(torch.cat([f1,f2],dim=1)))
        f4=self.DG3(self.conv1x1_2(torch.cat([f1,f2,f3],dim=1)))

        extract_f1=self.longRangeSkip(f1)
        extract_f2=self.conv1x1_3(torch.cat([f2,f3,f4],dim=1))

        h=self.add(extract_f1, extract_f2)
        h=self.OSMBlock(h)
        f_x4_1=F.upsample(h, size=[img.shape[2]*4, img.shape[3]*4], mode="bicubic")
        f_x4_2=F.upsample(img, size=[img.shape[2]*4, img.shape[3]*4], mode="bicubic")
        out=f_x4_1+f_x4_2

        out = self.add_mean(out)
        return out
//  [python]
def forward(self, img):

        img = self.sub_mean(img)

        f1=self.conv3x3(img)
        f2=self.DG1(f1)
        f3=self.DG2(self.conv1x1_1(torch.cat([f1,f2],dim=1)))
        f4=self.DG3(self.conv1x1_2(torch.cat([f1,f2,f3],dim=1)))

        extract_f1=self.longRangeSkip(f1)
        extract_f2=self.conv1x1_3(torch.cat([f2,f3,f4],dim=1))

        h=self.add(extract_f1, extract_f2)
        h=self.OSMBlock(h)
        f_x4_1=F.interpolate(h, size=[img.shape[2]*4, img.shape[3]*4], mode="bicubic",align_corners=False)
        f_x4_2=F.interpolate(img, size=[img.shape[2]*4, img.shape[3]*4], mode="bicubic",align_corners=False)
        out=f_x4_1+f_x4_2

        out = self.add_mean(out)
        return out
//  [python]
def forward(self, x): #l,m,s
        """
        # 128, 256, 512
        512, 256, 128
        from small -> large
        """
        x_level_0=x[2] #l
        x_level_1=x[1] #m
        x_level_2=x[0] #s
        # print('x_level_0: ', x_level_0.shape)
        # print('x_level_1: ', x_level_1.shape)
        # print('x_level_2: ', x_level_2.shape)
        if self.level == 0:
            level_0_resized = x_level_0
            level_1_resized = self.stride_level_1(x_level_1)
            level_2_downsampled_inter = F.max_pool2d(
                x_level_2, 3, stride=2, padding=1)
            level_2_resized = self.stride_level_2(level_2_downsampled_inter)
        elif self.level == 1:
            level_0_compressed = self.compress_level_0(x_level_0)
            level_0_resized = F.upsample(
                level_0_compressed, scale_factor=2, mode='nearest')
            level_1_resized = x_level_1
            level_2_resized = self.stride_level_2(x_level_2)
        elif self.level == 2:
            level_0_compressed = self.compress_level_0(x_level_0)
            level_0_resized = F.upsample(
                level_0_compressed, scale_factor=4, mode='nearest')
            x_level_1_compressed = self.compress_level_1(x_level_1)
            level_1_resized = F.upsample(
                x_level_1_compressed, scale_factor=2, mode='nearest')
            level_2_resized = x_level_2

        # print('level: {}, l1_resized: {}, l2_resized: {}'.format(self.level,
        #      level_1_resized.shape, level_2_resized.shape))
        level_0_weight_v = self.weight_level_0(level_0_resized)
        level_1_weight_v = self.weight_level_1(level_1_resized)
        level_2_weight_v = self.weight_level_2(level_2_resized)
        # print('level_0_weight_v: ', level_0_weight_v.shape)
        # print('level_1_weight_v: ', level_1_weight_v.shape)
        # print('level_2_weight_v: ', level_2_weight_v.shape)

        levels_weight_v = torch.cat(
            (level_0_weight_v, level_1_weight_v, level_2_weight_v), 1)
        levels_weight = self.weight_levels(levels_weight_v)
        levels_weight = F.softmax(levels_weight, dim=1)

        fused_out_reduced = level_0_resized * levels_weight[:, 0:1, :, :] +\
            level_1_resized * levels_weight[:, 1:2, :, :] +\
            level_2_resized * levels_weight[:, 2:, :, :]

        out = self.expand(fused_out_reduced)

        if self.vis:
            return out, levels_weight, fused_out_reduced.sum(dim=1)
        else:
            return out
//  [python]
def forward(self, x): #l,m,s
        """
        # 128, 256, 512
        512, 256, 128
        from small -> large
        """
        x_level_0=x[2] #l
        x_level_1=x[1] #m
        x_level_2=x[0] #s
        # print('x_level_0: ', x_level_0.shape)
        # print('x_level_1: ', x_level_1.shape)
        # print('x_level_2: ', x_level_2.shape)
        if self.level == 0:
            level_0_resized = x_level_0
            level_1_resized = self.stride_level_1(x_level_1)
            level_2_downsampled_inter = F.max_pool2d(
                x_level_2, 3, stride=2, padding=1)
            level_2_resized = self.stride_level_2(level_2_downsampled_inter)
        elif self.level == 1:
            level_0_compressed = self.compress_level_0(x_level_0)
            level_0_resized = F.interpolate(
                level_0_compressed, scale_factor=2, mode='nearest')
            level_1_resized = x_level_1
            level_2_resized = self.stride_level_2(x_level_2)
        elif self.level == 2:
            level_0_compressed = self.compress_level_0(x_level_0)
            level_0_resized = F.interpolate(
                level_0_compressed, scale_factor=4, mode='nearest')
            x_level_1_compressed = self.compress_level_1(x_level_1)
            level_1_resized = F.interpolate(
                x_level_1_compressed, scale_factor=2, mode='nearest')
            level_2_resized = x_level_2

        # print('level: {}, l1_resized: {}, l2_resized: {}'.format(self.level,
        #      level_1_resized.shape, level_2_resized.shape))
        level_0_weight_v = self.weight_level_0(level_0_resized)
        level_1_weight_v = self.weight_level_1(level_1_resized)
        level_2_weight_v = self.weight_level_2(level_2_resized)
        # print('level_0_weight_v: ', level_0_weight_v.shape)
        # print('level_1_weight_v: ', level_1_weight_v.shape)
        # print('level_2_weight_v: ', level_2_weight_v.shape)

        levels_weight_v = torch.cat(
            (level_0_weight_v, level_1_weight_v, level_2_weight_v), 1)
        levels_weight = self.weight_levels(levels_weight_v)
        levels_weight = F.softmax(levels_weight, dim=1)

        fused_out_reduced = level_0_resized * levels_weight[:, 0:1, :, :] +\
            level_1_resized * levels_weight[:, 1:2, :, :] +\
            level_2_resized * levels_weight[:, 2:, :, :]

        out = self.expand(fused_out_reduced)

        if self.vis:
            return out, levels_weight, fused_out_reduced.sum(dim=1)
        else:
            return out
//  [python]
def forward(self, x):
        z = self.model.encode(x).sample()
        z = z * self.scale_factor
        noise_level = torch.randint(0, self.max_noise_level, (x.shape[0],), device=x.device).long()
        z = self.q_sample(z, noise_level)
        if self.out_size is not None:
            z = t.interpolate(z, size=self.out_size, mode="nearest")  # TODO: experiment with mode
        # z = z.repeat_interleave(2, -2).repeat_interleave(2, -1)
        return z, noise_level
//  [python]
def forward(self, x):
        z = self.model.encode(x).sample()
        z = z * self.scale_factor
        noise_level = torch.randint(0, self.max_noise_level, (x.shape[0],), device=x.device).long()
        z = self.q_sample(z, noise_level)
        if self.out_size is not None:
            z = torch.nn.functional.interpolate(z, size=self.out_size, mode="nearest")  # TODO: experiment with mode
        # z = z.repeat_interleave(2, -2).repeat_interleave(2, -1)
        return z, noise_level
//  [python]
def forward(self, x, position_embedding=None, **kwargs):
        # FIXME look at relaxing size constraints
        # assert H == self.img_size[0] and W == self.img_size[1], \
        #     f"Input image size ({H}*{W}) doesn't match model ({self.img_size[0]}*{self.img_size[1]})."
        x = self.proj(x)
        Hp, Wp = x.shape[2], x.shape[3]

        if position_embedding is not None:
            # interpolate the position embedding to the corresponding size
            position_embedding = position_embedding.view(1, self.patch_shape[0], self.patch_shape[1], -1).permute(0, 3,
                                                                                                                  1, 2)
            position_embedding = t.interpolate(position_embedding, size=(Hp, Wp), mode='bicubic')
            x = x + position_embedding

        x = x.flatten(2).transpose(1, 2)
        return x, (Hp, Wp)
//  [python]
def forward(self, x, position_embedding=None, **kwargs):
        # FIXME look at relaxing size constraints
        # assert H == self.img_size[0] and W == self.img_size[1], \
        #     f"Input image size ({H}*{W}) doesn't match model ({self.img_size[0]}*{self.img_size[1]})."
        x = self.proj(x)
        Hp, Wp = x.shape[2], x.shape[3]

        if position_embedding is not None:
            # interpolate the position embedding to the corresponding size
            position_embedding = position_embedding.view(1, self.patch_shape[0], self.patch_shape[1], -1).permute(0, 3,
                                                                                                                  1, 2)
            position_embedding = F.interpolate(position_embedding, size=(Hp, Wp), mode='bicubic')
            x = x + position_embedding

        x = x.flatten(2).transpose(1, 2)
        return x, (Hp, Wp)
//  [python]
def forward(self, ref, supp):
        """
        Args:
            ref (Tensor): Reference image with shape of (b, 3, h, w).
            supp: The supporting image to be warped: (b, 3, h, w).

        Returns:
            Tensor: Estimated optical flow: (b, 2, h, w).
        """
        num_batches, _, h, w = ref.size()
        ref = [ref]
        supp = [supp]

        # generate downsampled frames
        for _ in range(3):
            ref.insert(
                0,
                F.avg_pool2d(
                    input=ref[0],
                    kernel_size=2,
                    stride=2,
                    count_include_pad=False))
            supp.insert(
                0,
                F.avg_pool2d(
                    input=supp[0],
                    kernel_size=2,
                    stride=2,
                    count_include_pad=False))

        # flow computation
        flow = ref[0].new_zeros(num_batches, 2, h // 16, w // 16)
        for i in range(4):
            flow_up = t.upsample(
                input=flow,
                scale_factor=2,
                mode='bilinear') * 2.0
            flow = flow_up + self.basic_module[i](
                torch.cat([
                    ref[i],
                    flow_warp(supp[i], flow_up.permute(0, 2, 3, 1)), flow_up
                ], 1))
        return flow
//  [python]
def forward(self, ref, supp):
        """
        Args:
            ref (Tensor): Reference image with shape of (b, 3, h, w).
            supp: The supporting image to be warped: (b, 3, h, w).

        Returns:
            Tensor: Estimated optical flow: (b, 2, h, w).
        """
        num_batches, _, h, w = ref.size()
        ref = [ref]
        supp = [supp]

        # generate downsampled frames
        for _ in range(3):
            ref.insert(
                0,
                F.avg_pool2d(
                    input=ref[0],
                    kernel_size=2,
                    stride=2,
                    count_include_pad=False))
            supp.insert(
                0,
                F.avg_pool2d(
                    input=supp[0],
                    kernel_size=2,
                    stride=2,
                    count_include_pad=False))

        # flow computation
        flow = ref[0].new_zeros(num_batches, 2, h // 16, w // 16)
        for i in range(4):
            flow_up = F.interpolate(
                input=flow,
                scale_factor=2,
                mode='bilinear',
                align_corners=True) * 2.0
            flow = flow_up + self.basic_module[i](
                torch.cat([
                    ref[i],
                    flow_warp(supp[i], flow_up.permute(0, 2, 3, 1)), flow_up
                ], 1))
        return flow