text: """
  Take care of SEP case.
  """
  
  from mdl_conc_single import ConcBase
  import torch
  from torch import nn
  from torch.nn import functional as F
  from mdl_srl_utils import combine_first_ax
  from box_utils import bbox_overlaps
  
  
  class ConcSEP(ConcBase):
      def conc_encode(self, conc_feats, inp):
          nfrm = self.num_sampled_frm
          nppf = self.num_prop_per_frm
          ncmp = inp['new_srl_idxs'].size(1)
          return self.conc_encode_item(conc_feats, inp, nfrm, nppf, ncmp)
  
      def simple_obj_interact_input(self, prop_seg_feats, inp):
          B, num_cmp, num_props, psdim = prop_seg_feats.shape
          return self.simple_obj_interact(
              prop_seg_feats, inp,
              num_cmp, self.num_sampled_frm,
              self.num_prop_per_frm
          )
  
      def set_args_conc(self):
          self.nfrms = self.num_sampled_frm
          self.nppf = self.num_prop_per_frm
  
      def get_num_cmp_msk(self, inp, out_shape):
          num_cmp = inp['new_srl_idxs'].size(1)
          B, num_verbs, num_srl_args, seq_len = inp['srl_arg_words_ind'].shape
          num_cmp_msk = inp['num_cmp_msk'].view(
              B, num_cmp, 1, 1
          ).expand(
              B, num_cmp, num_srl_args,
              self.num_sampled_frm * self.num_prop_per_frm
          ).contiguous(
          ).view(*out_shape)
          return num_cmp_msk
  
      def concat_prop_seg_feats(self, prop_feats, seg_feats, inp):
          B, num_cmp, num_props, pdim = prop_feats.shape
          prop_seg_feats = torch.cat(
              [
                  prop_feats.view(
                      B, num_cmp, self.num_sampled_frm,
                      self.num_prop_per_frm, prop_feats.size(-1)
                  ),
                  seg_feats.unsqueeze(-2).expand(
                      B, num_cmp, self.num_sampled_frm,
                      self.num_prop_per_frm, seg_feats.size(-1)
                  )
              ], dim=-1
          ).view(
              B, num_cmp, self.num_sampled_frm*self.num_prop_per_frm,
              prop_feats.size(-1) + seg_feats.size(-1)
          )
          # B x num_cmp x nfrm*nppf x psdim
          return prop_seg_feats
  
      def compute_fin_scores(self, conc_out_dict, inp, vidf_outs=None):
          """
          output fin scores should be of shape
          B x num_cmp
          prop_scores: B x num_cmp x num_srl_args x num_props
          """
          prop_scores1 = conc_out_dict['conc_feats_out'].clone().detach()
          prop_scores = torch.sigmoid(prop_scores1)
          # prop_scores = prop_scores1
          if self.cfg.mdl.use_vis_msk:
              # B x num_cmp x num_srl_args
              prop_scores_max_boxes, _ = torch.max(prop_scores, dim=-1)
  
              # B x num_cmp x num_srl_args
              srl_arg_inds_msk = inp['srl_arg_inds_msk'].float()
              B, num_verbs, num_srl_args = srl_arg_inds_msk.shape
  
              num_cmp = prop_scores.size(1)
  
              if vidf_outs is not None:
                  # add vidf outs to the verb places
                  vidf_outs = torch.sigmoid(vidf_outs)
  
                  # B x num_cmp -> B x num_cmp x num_srl_args
                  vidf_outs = vidf_outs.unsqueeze(-1).expand(
                      *prop_scores_max_boxes.shape
                  )
                  vmsk = inp['verb_ind_in_srl']
  
                  if vmsk.size(1) == 1 and num_cmp > 1:
                      vmsk = vmsk.expand(-1, num_cmp)
                  # B x num_cmp
                  vmsk = vmsk.view(
                      B, num_cmp, 1).expand(
                          B, num_cmp, num_srl_args
                  )
                  prop_scores_max_boxes.scatter_(
                      dim=2,
                      index=vmsk,
                      src=vidf_outs
                  )
  
              prop_scores_max_boxes = prop_scores_max_boxes * srl_arg_inds_msk
  
              # b x num_cmp
              fin_scores_eval = prop_scores_max_boxes.sum(
                  dim=-1) / srl_arg_inds_msk.sum(dim=-1)
  
              verb_msk = inp['num_cmp_msk']
              fin_scores_eval = fin_scores_eval * verb_msk.float()
  
              fin_scores_loss = prop_scores_max_boxes * verb_msk.unsqueeze(
                  -1).expand(*prop_scores_max_boxes.shape).float()
              return {
                  # B x num_cmp
                  'fin_scores_eval': fin_scores_eval,
                  # B x num_cmp x num_srl_args
                  'fin_scores_loss': fin_scores_loss
              }
  
          else:
              # B x num_cmp x num_cmp x num_srl_args
              prop_scores_max_boxes, _ = torch.max(prop_scores, dim=-1)
              # B x num_cmp x num_cmp
              fin_scores = prop_scores_max_boxes.sum(dim=-1)
          return fin_scores
  
      def forward(self, inp):
          """
          Main difference is that prop feats/seg features
          have an extra dimension
          """
          # B x 6 x 5 x 40
          # 6 is num_cmp for a sent
          # 5 is num args in a sent
          # 40 is seq length for each arg
          B, num_verbs, num_srl_args, seq_len = inp['srl_arg_words_ind'].shape
          # B*num_cmp x seq_len
          src_toks = self.get_srl_arg_seq_to_sent_seq(inp)
          # B*num_cmp x seq_len
          src_lens = inp['srl_arg_word_mask_len'].view(B*num_verbs, -1)
          # B*num_cmp x seq_len x 256
          lstm_outs = self.lang_encode(src_toks, src_lens)
          lstm_encoded = lstm_outs['lstm_full_output']
  
          # B x num_cmp x 5 x 512
          srl_arg_lstm_encoded = self.retrieve_srl_arg_from_lang_encode(
              lstm_encoded, inp
          )
  
          # Get visual features
          # B x num_cmp x 1000 x 512
          prop_feats = self.prop_feats_encode(inp)
          # B, num_cmp, num_props, pdim = prop_feats.shape
  
          # Get seg features
          # B x num_cmp x 10 x 512
          seg_feats = self.seg_feats_encode(inp)
  
          # B x num_cmp x nfrm*nppf x psdim
          prop_seg_feats = self.concat_prop_seg_feats(prop_feats, seg_feats, inp)
  
          prop_seg_feats = self.simple_obj_interact_input(
              prop_seg_feats, inp
          )
  
          num_cmp = inp['new_srl_idxs'].size(1)
          if srl_arg_lstm_encoded.size(1) == 1 and num_cmp > 1:
              srl_arg_lstm_encoded = srl_arg_lstm_encoded.expand(
                  -1, num_cmp, -1, -1
              )
  
          conc_feats = self.concate_vis_lang_feats(
              prop_seg_feats, srl_arg_lstm_encoded
          )
  
          # B x num_cmp x num_srl_args x num_props
          conc_feats_out_dict = self.conc_encode(conc_feats, inp)
          conc_feats_out = conc_feats_out_dict['conc_feats_out']
  
          seg_feats_for_verb, verb_feats = self.get_seg_verb_feats_to_process(
              seg_feats, srl_arg_lstm_encoded, lstm_outs, inp
          )
  
          if verb_feats.size(1) == 1 and num_cmp > 1:
              verb_feats = verb_feats.expand(-1, num_cmp, -1)
  
          # B x num_cmp
          vidf_outs = self.compute_seg_verb_feats_out(
              seg_feats_for_verb, verb_feats
          )
          fin_scores = self.compute_fin_scores(
              conc_feats_out_dict, inp, vidf_outs
          )
  
          num_cmp_msk = self.get_num_cmp_msk(inp, conc_feats_out.shape)
  
          srl_ind_msk = inp['srl_arg_inds_msk']
          if srl_ind_msk.size(1) == 1 and num_cmp > 1:
              srl_ind_msk = srl_ind_msk.expand(
                  -1, num_cmp, -1, -1
              )
          srl_ind_msk = srl_ind_msk.unsqueeze(-1).expand(
              *conc_feats_out.shape)
          mdl_outs_eval = torch.sigmoid(
              conc_feats_out) * srl_ind_msk.float() * num_cmp_msk.float()
  
          return {
              'mdl_outs': conc_feats_out,
              'mdl_outs_eval': mdl_outs_eval,
              'vidf_outs': vidf_outs,
              'fin_scores_loss': fin_scores['fin_scores_loss'],
              'fin_scores': fin_scores['fin_scores_eval']
          }
  
  
  class LossB_SEP(nn.Module):
      """
      Loss Function (for a batch) for SEP case.
      Specifically, we need to have a separate verb loss
      Also, handling of some functions is different
      from single video case
      """
  
      def __init__(self, cfg, comm):
          super().__init__()
          self.cfg = cfg
          self.comm = comm
          self.loss_keys = ['loss', 'mdl_out_loss', 'verb_loss']
          self.loss_lambda = self.cfg.loss.loss_lambda
          self.after_init()
  
      def after_init(self):
          pass
  
      def get_targets_from_overlaps(self, overlaps, inp):
          """
          Use the given overlaps to produce the targets
          overlaps: B x num_cmp x 1000 x 100
          """
          targets = overlaps
  
          srl_boxes = inp['srl_boxes']
          B, num_verbs, num_srl_args, num_box_per_srl = srl_boxes.shape
          B, num_cmp, num_props, num_gt_box = targets.shape
  
          if num_verbs == 1 and num_cmp > 1:
              srl_boxes = srl_boxes.expand(-1, num_cmp, -1, -1)
  
          srl_boxes_reshaped = srl_boxes.view(
              B, num_cmp, num_srl_args, 1, num_box_per_srl).expand(
                  B, num_cmp, num_srl_args, num_props, num_box_per_srl)
  
          targets_reshaped = targets.view(
              B, num_cmp, 1, num_props, num_gt_box).expand(
                  B, num_cmp, num_srl_args, num_props, num_gt_box)
  
          # Choose only those proposals which are ground-truth
          # for given srl
          targets_to_use = torch.gather(
              targets_reshaped, dim=-1, index=srl_boxes_reshaped)
  
          srl_boxes_lens = inp['srl_boxes_lens']
          targets_to_use = (
              targets_to_use * srl_boxes_lens.float().unsqueeze(
                  -2).expand(*targets_to_use.shape)
          )
  
          targets_to_use = targets_to_use.max(dim=-1)[0] > 0.5
  
          return targets_to_use
  
      def compute_overlaps(self, inp):
  
          pad_props = inp['pad_proposals']
          gt_bboxs = inp['pad_gt_bboxs']
          frm_msk = inp['pad_frm_mask']
          pnt_msk = inp['pad_pnt_mask']
  
          assert len(pnt_msk.shape) == 3
  
          B = pad_props.size(0)
          num_cmp = pad_props.size(1)
          pad_props = combine_first_ax(pad_props)
          gt_bboxs = combine_first_ax(gt_bboxs)
          frm_msk = combine_first_ax(frm_msk)
  
          pnt_msk = combine_first_ax(pnt_msk)
  
          overlaps = bbox_overlaps(
              pad_props, gt_bboxs,
              (frm_msk | pnt_msk[:, :].unsqueeze(-1)))
          overlaps = overlaps.view(B, num_cmp, *overlaps.shape[1:])
  
          return overlaps
  
      def compute_loss_targets(self, inp):
          """
          Compute the targets, based on iou
          overlaps
          """
          overlaps = self.compute_overlaps(inp)
          B, ncmp, nprop, ngt = overlaps.shape
          overlaps_msk = overlaps.new_zeros(*overlaps.shape)
  
          targ_cmp = inp['target_cmp']
          # overlaps_msk[:, targ_cmp, ...] = 1
          overlaps_msk.scatter_(
              dim=1,
              index=targ_cmp.view(B, 1, 1, 1).expand(B, ncmp, nprop, ngt),
              src=overlaps_msk.new_ones(*overlaps_msk.shape)
          )
  
          overlaps_one_targ = overlaps * overlaps_msk
  
          targets_one = self.get_targets_from_overlaps(overlaps_one_targ, inp)
          targets_all = self.get_targets_from_overlaps(overlaps, inp)
          return {
              'targets_one': targets_one,
              'targets_all': targets_all
          }
  
      def compute_mdl_loss(self, mdl_outs, targets_one, inp):
          weights = None
          tot_loss = F.binary_cross_entropy_with_logits(
              mdl_outs, target=targets_one.float(),
              weight=weights,
              reduction='none'
          )
  
          # B x num_cmp
          num_cmp_msk = inp['num_cmp_msk']
          num_cmp = num_cmp_msk.size(1)
          srl_arg_boxes_mask = inp['srl_arg_boxes_mask']
          num_verbs = srl_arg_boxes_mask.size(1)
          if num_verbs == 1 and num_cmp > 1:
              srl_arg_boxes_mask = srl_arg_boxes_mask.expand(-1, num_cmp, -1)
  
          B, num_cmp, num_srl_args = srl_arg_boxes_mask.shape
  
          boxes_msk = num_cmp_msk.unsqueeze(
              -1).expand(*srl_arg_boxes_mask.shape).float()
  
          # B x num_cmp x num_srl_args -> B x num_cmp x num_srl x 1000
          boxes_msk = boxes_msk.unsqueeze(
              -1).expand(*targets_one.shape)
  
          tot_loss = tot_loss * boxes_msk
  
          multiplier = tot_loss.size(-1)
          if srl_arg_boxes_mask.max() > 0:
              out_loss = torch.masked_select(tot_loss, boxes_msk.byte())
          else:
              # TODO: NEED TO check what is wrong here
              out_loss = tot_loss
  
          mdl_out_loss = out_loss.mean() * multiplier
  
          return mdl_out_loss
  
      def compute_vidf_loss_simple(self, vidf_outs, inp):
          """
          vidf_outs are fin scores: B x ncmp x nfrms
          """
          B, ncmp, nfrm = vidf_outs.shape
          targs = vidf_outs.new_zeros(*vidf_outs.shape)
  
          targ_cmp = inp['target_cmp']
  
          targs.scatter_(
              dim=1,
              index=targ_cmp.view(B, 1, 1).expand(B, ncmp, nfrm),
              src=targs.new_ones(*targs.shape)
          )
  
          # B x ncmp x nfrms
          out_loss = F.binary_cross_entropy(vidf_outs, targs, reduction='none')
  
          mult = 1. / nfrm
  
          # B x ncmp
          msk = inp['num_cmp_msk']
          out_loss = torch.masked_select(out_loss.sum(dim=-1) * msk.float(),
                                         msk.byte()) * mult
          return out_loss.mean()
  
      def compute_vidf_loss(self, vidf_outs, inp):
          B, num_cmp, num_srl_args = vidf_outs.shape
          box_msk = inp['srl_arg_boxes_mask']
          srl_arg_ind_msk = inp['srl_arg_inds_msk']
          vidf_outs = ((vidf_outs * box_msk.float()).sum(dim=-1) /
                       srl_arg_ind_msk.sum(dim=-1).float())
          vidf_targs = vidf_outs.new_zeros(*vidf_outs.shape)
  
          targ_cmp = inp['target_cmp']
  
          vidf_targs.scatter_(
              dim=1,
              index=targ_cmp.unsqueeze(-1).expand(*vidf_targs.shape),
              src=vidf_targs.new_ones(*vidf_targs.shape)
          )
  
          vidf_loss = F.binary_cross_entropy(  #
              vidf_outs, vidf_targs,
              reduction='none'
          )
          msk = inp['num_cmp_msk']
          vidf_loss = vidf_loss * msk.float()
          vidf_loss = torch.masked_select(vidf_loss, msk.byte())
          return vidf_loss.mean()
  
      def forward(self, out, inp):
          targets_all = self.compute_loss_targets(inp)
          targets_n = targets_all['targets_one']
  
          mdl_outs = out['mdl_outs']
  
          mdl_out_loss = self.compute_mdl_loss(mdl_outs, targets_n, inp)
  
          verb_outs = out['vidf_outs']
  
          verb_loss = F.binary_cross_entropy_with_logits(
              verb_outs,
              inp['verb_cmp'].float(),
              reduction='none'
          )
  
          vcc_msk = inp['verb_cross_cmp_msk'].float()
          vcc_msk = (vcc_msk.sum(dim=-1) > 0).float()
  
          verb_loss = verb_loss * vcc_msk
          verb_loss = torch.masked_select(
              verb_loss, vcc_msk.byte()).mean()
  
          # out_loss = mdl_out_loss + verb_loss
          out_loss = mdl_out_loss
  
          out_loss_dict = {
              'loss': out_loss,
              'mdl_out_loss': mdl_out_loss,
              'verb_loss': verb_loss
          }
  
          return {k: v * self.loss_lambda for k, v in out_loss_dict.items()}
  