Ejemplo n.º 1
0
    def process(self, images, return_time=False):

        with torch.no_grad():
            output = self.model(images)[-1]
            hm_obj = output['hm'].sigmoid_()
            hm_human = output['hm_human'].sigmoid_()
            hm_rel = output['hm_rel'].sigmoid_()
            wh = output['wh']
            reg = output['reg'] if self.opt.reg_offset else None
            sub_offset = output['sub_offset']
            obj_offset = output['obj_offset']
            torch.cuda.synchronize()
            forward_time = time.time()

            dets_obj, dets_sub, rel = hoidet_decode(
                hm_human,
                hm_obj,
                wh,
                hm_rel,
                sub_offset,
                obj_offset,
                reg=reg,
                corremat=self.corre_mat,
                is_sub_verb=self.opt.use_verb_sub)

        if return_time:
            return output, dets_obj, dets_sub, rel, forward_time, images.size(
            )[2], images.size()[3]
        else:
            return output, dets_obj, dets_sub, rel, images.size(
            )[2], images.size()[3]
Ejemplo n.º 2
0
    def process(self, images, return_time=False):

        with torch.no_grad():
            output = self.model(images)[-1]
            hm_obj = output['hm'].sigmoid_()

            if self.opt.test_mul:
                hm_rel = output['hm_rel_refine1'].sigmoid_() * output[
                    'hm_rel_refine'].sigmoid_() * output['hm_rel'].sigmoid_()
            elif self.opt.test_add:
                hm_rel = (output['hm_rel_refine_1'] * 0.1 +
                          output['hm_rel'] * 0.2 +
                          output['hm_rel_refine']).sigmoid_()
            elif self.opt.test_refine:
                hm_rel = output['hm_rel_refine'].sigmoid_()
            else:
                hm_rel = output['hm_rel'].sigmoid_()

            wh = output['wh']
            reg = output['reg'] if self.opt.reg_offset else None

            sub_offset = output['sub_offset']
            obj_offset = output['obj_offset']
            torch.cuda.synchronize()
            forward_time = time.time()
            if self.opt.hm_rel_dcn25_i_casc_match or self.opt.hm_rel_dcn25_i_match:
                dets_obj, dets_sub, rel = hoidet_decode_match(
                    hm_obj,
                    wh,
                    hm_rel,
                    sub_offset,
                    obj_offset,
                    reg=reg,
                    corremat=self.corre_mat,
                    is_sub_verb=self.opt.use_verb_sub,
                    K_rel=self.opt.test_top_K_rel)
            else:
                dets_obj, dets_sub, rel = hoidet_decode(
                    hm_obj,
                    wh,
                    hm_rel,
                    sub_offset,
                    obj_offset,
                    reg=reg,
                    corremat=self.corre_mat,
                    is_sub_verb=self.opt.use_verb_sub,
                    K_rel=self.opt.test_top_K_rel)

        if return_time:
            return output, dets_obj, dets_sub, rel, forward_time, images.size(
            )[2], images.size()[3]
        else:
            return output, dets_obj, dets_sub, rel, images.size(
            )[2], images.size()[3]