def loss_spatials(self, outputs, targets, indices, num_masks, layer_id, extra):
        """Compute the losses related to the masks: the focal loss and the dice loss.
        targets dicts must contain the key "masks" containing a tensor of dim [nb_target_boxes, h, w]
        """
        assert "pred_smasks" in outputs
        assert "pred_smaskembs" in outputs

        if layer_id >= self.top_x_layers['spatial']:
            loss = outputs['pred_smasks'].sum() * 0.0
            loss_grd_ce = outputs["pred_smasks"].sum() * 0.0
            return {"loss_spatial_bce_0": loss, "loss_spatial_dice_0": loss, "loss_spatial_ce_0": loss_grd_ce}

        gt_masks = [x['gt_spatial_masks'] for x in targets]
        # compute a keep index with batch size to avoid empty gt_masks
        stack_gt_mask = torch.cat(gt_masks)
        bs,_,_ = stack_gt_mask.shape
        stack_gt_mask = stack_gt_mask.view(bs,-1).sum(dim=-1)
        keep = stack_gt_mask > 0 # only keep sample contain positive mask

        if keep.sum() == 0:
            loss = outputs['pred_smasks'].sum() * 0.0
            loss_grd_ce = outputs["pred_smasks"].sum() * 0.0
            return {"loss_spatial_bce_0": loss, "loss_spatial_dice_0": loss, "loss_spatial_ce_0": loss_grd_ce}

        # mask embedding logits
        v_emb = outputs["pred_smaskembs"] # [bs, nq, 512]

        # pos mask
        s_emb = outputs["pred_pspatials"] # [bs, ns, 512]
        pred_logits = v_emb @ s_emb.transpose(1,2)
        outputs['pred_pos_logits'] = pred_logits # [bs, nq, 1]
        indices = self.matcher(outputs, targets, mode='spatial', extra={})
        src_idx = self._get_src_permutation_idx(indices)
        tgt_idx = self._get_tgt_permutation_idx(indices)

        # pos class loss
        pred_logit = torch.cat([o[:len(t['gt_spatial_masks'])] for o,t in zip(outputs["pred_pos_logits"].transpose(1,2), targets)])
        gt_logit = torch.zeros_like(pred_logit)
        gt_logit = gt_logit[keep]
        _src_idx = [torch.arange(keep.sum(), device=src_idx[0].device), src_idx[1][keep.cpu()]]
        gt_logit[_src_idx] = 1
        pred_logit = pred_logit[keep]
        loss_spa_ce_pos = torch.sum(-gt_logit * F.log_softmax(pred_logit, dim=-1), dim=-1).mean()

        # neg mask
        # s_emb = outputs["pred_nspatials"] # [bs, ns, 512]
        # neg_mask = (s_emb.sum(dim=list(range(1, len(s_emb.shape)))) != 0).float()[keep]
        # pred_logits = v_emb @ s_emb.transpose(1,2)
        # outputs['pred_neg_logits'] = pred_logits # [bs, nq, 1]
        # indices = self.matcher(outputs, targets, mode='spatial_pn', extra=extra)
        # src_idx = self._get_src_permutation_idx(indices)
        # tgt_idx = self._get_tgt_permutation_idx(indices)
        # src_masks_neg = outputs["pred_smasks"][src_idx][keep]
        # src_masks_neg = src_masks_neg*(neg_mask[:,None,None])
        # src_masks_neg = src_masks_neg.clip(0) * (-1)

        # neg class loss
        # pred_logit = outputs["pred_neg_logits"]
        # gt_logit = torch.zeros_like(pred_logit)
        # gt_logit[src_idx] = 1
        # bs,_,ns = pred_logit[keep].shape
        # pred_logit = pred_logit[keep].transpose(1,2).view(bs*ns,-1)
        # gt_logit = gt_logit[keep].transpose(1,2).view(bs*ns,-1)
        # loss_spa_ce_neg = (torch.sum(-gt_logit * F.log_softmax(pred_logit, dim=-1), dim=-1)*neg_mask).sum() / (neg_mask.sum()+1e-6)

        # recompute a keep index with matched tgt
        stack_gt_mask = nn.utils.rnn.pad_sequence(gt_masks, padding_value=-1).transpose(0,1)[tgt_idx]        
        bs,_,_ = stack_gt_mask.shape
        target_masks = stack_gt_mask
        stack_gt_mask = stack_gt_mask.view(bs,-1).sum(dim=-1)
        keep = stack_gt_mask > 0 # only keep sample contain positive mask
        src_masks_pos = outputs["pred_smasks"][src_idx][keep]

        # TODO use valid to mask invalid areas due to padding in loss
        target_masks = target_masks.to(src_masks_pos)
        target_masks = target_masks[keep]

        # mul = extra['spatial_query_mode'][keep]
        # src_masks_cur = src_masks_cur.clip(0) * mul[:,None,None]
        # src_masks_cur = src_masks_cur

        # if neg_mask[0] == 1:
        #     import cv2
        #     print(src_masks_pos.shape)
        #     print(src_masks_neg.shape)
        #     print(target_masks.shape)
        #     # import pdb; pdb.set_trace()
        #     v_pos_mask = (src_masks_pos[0].sigmoid() > 0.5).float().cpu().detach().numpy() * 255
        #     v_neg_mask = (_src_masks_neg[0].sigmoid() > 0.5).float().cpu().detach().numpy() * 255
        #     v_sum = ((src_masks_pos[0]-_src_masks_neg[0].clip(0)).sigmoid() > 0.5).float().cpu().detach().numpy() * 255
        #     v_gt = target_masks[0].float().cpu().detach().numpy() * 255

        #     cv2.imwrite('v_pos_mask.png', v_pos_mask)
        #     cv2.imwrite('v_neg_mask.png', v_neg_mask)
        #     cv2.imwrite('v_sum.png', v_sum)
        #     cv2.imwrite('v_gt.png', v_gt)
        #     import pdb; pdb.set_trace()

        # src_masks = (src_masks_pos + src_masks_neg)[:, None]
        src_masks = src_masks_pos[:, None]
        target_masks = target_masks[:, None]

        # debug visualization
        # with torch.no_grad():
        #     import cv2
        #     import numpy as np

        #     v_src_masks = (F.interpolate(src_masks, size=target_masks.shape[-2:], mode='bilinear', align_corners=False).sigmoid() > 0.5).float().cpu().numpy()[:,0] * 255
        #     v_target_masks = target_masks.float().cpu().numpy()[:,0] * 255
        #     v_masks = np.concatenate([v_src_masks, v_target_masks], axis=2)

        #     for i in range(len(src_masks)):
        #         v1 = v_src_masks[i]
        #         v2 = v_target_masks[i]
        #         v = np.concatenate([v1,v2], axis=1)
        #         cv2.imwrite('v{}.png'.format(i), v)
        #     import pdb; pdb.set_trace()

        # visualization
        # VL.step()
        # v_img = batched_inputs[0]['image'].permute(1,2,0).cpu().numpy()
        # VL.add_image(v_img[:,:,::-1])
        # candidate_masks = batched_inputs[0]['spatial_query']['rand_shape'].float().cpu().numpy()
        # gt_masks = batched_inputs[0]['spatial_query']['gt_masks'].float().cpu().numpy()
        # texts = ['cmask' for i in range(len(candidate_masks))]
        # VL.overlay_obj_mask_to_image(v_img[:,:,::-1], candidate_masks, texts)
        # texts = ['gmask' for i in range(len(candidate_masks))]
        # VL.overlay_obj_mask_to_image(v_img[:,:,::-1], gt_masks, texts)

        # import cv2
        # for i in range(len(src_masks)):
        #     visual_src_mask_cur = (src_masks_cur[i].sigmoid()>0.5).detach().float().cpu().numpy() * 255
        #     visual_src_mask_mem = (src_masks_mem[i].sigmoid()>0.5).detach().float().cpu().numpy() * 255
        #     visual_src_mask = (src_masks[i,0].sigmoid()>0.5).detach().float().cpu().numpy() * 255
        #     visual_target_mask = (target_masks[i,0].sigmoid()>0.5).detach().float().cpu().numpy() * 255

        #     cv2.imwrite('visual_src_mask_cur_{}_{}.png'.format(i, mul[i].item()), visual_src_mask_cur)
        #     cv2.imwrite('visual_src_mask_mem_{}_{}.png'.format(i, mul[i].item()), visual_src_mask_mem)
        #     cv2.imwrite('visual_src_mask_{}_{}.png'.format(i, mul[i].item()), visual_src_mask)
        #     cv2.imwrite('visual_target_mask_{}_{}.png'.format(i, mul[i].item()), visual_target_mask)
        # import pdb; pdb.set_trace()

        with torch.no_grad():
            # sample point_coords
            point_coords = get_uncertain_point_coords_with_randomness(
                src_masks,
                lambda logits: calculate_uncertainty(logits),
                self.num_points,
                self.oversample_ratio,
                self.importance_sample_ratio,
            ).type(src_masks.dtype)
            # get gt labels
            point_labels = point_sample(
                target_masks,
                point_coords,
                align_corners=False,
            ).squeeze(1)

        point_logits = point_sample(
            src_masks,
            point_coords,
            align_corners=False,
        ).squeeze(1)

        num_masks = len(src_masks)
        losses = {
            "loss_spatial_bce_0": sigmoid_ce_loss_jit(point_logits, point_labels, num_masks),
            "loss_spatial_dice_0": dice_loss_jit(point_logits, point_labels, num_masks),
        }

        # losses.update({"loss_spatial_ce_0": loss_spa_ce_pos + loss_spa_ce_neg})
        losses.update({"loss_spatial_ce_0": loss_spa_ce_pos})

        del src_masks
        del target_masks
        return losses
Compute the losses related to the masks: the focal loss and the dice loss.
        targets dicts must contain the key "masks" containing a tensor of dim [nb_target_boxes, h, w]
        
modeling/modules/criterion.py