Esempio n. 1
0
def more_actors_gurobi(data, n_actors, constraints,
                       first_run=False):
    """Multi-actor labeling using confidence weighted screen space distance.

    Args:
        first_run (bool):
            Short first run for vis only.
        n_actors (int):
            How many actors to label.
        constraints (Dict[str, Dict[int, int]]):
            {frame_str => {pose_id => actor_id}}.
        first_run (bool):
            Is this the very first run (limit runtime for only vis).
    Returns:
        pose_ids (Dict[int, Dict[int, int]]):
            {frame_id => {actor_id => pose_id}}
        problem (ActorProblem2):
            Labeling problem.
        data (PosesWrapper):
            Wrapped data for visualization.
    """

    # color_norm = cmNormalize(vmin=0, vmax=n_actors+1)
    # scalar_map = cm.ScalarMappable(norm=color_norm, cmap='gist_earth')
    # colors = [tuple(c * 255. for c in scalar_map.to_rgba(i+1))
    #           for i in range(n_actors)]
    # print(colors)
    # raise RuntimeError("")

    if isinstance(data, Skeleton):
        data = SkeletonPosesWrapper(skeleton=data)
    else:
        assert isinstance(data, dict), "%s" % type(data)
        data = DataPosesWrapper(data=data)
    is_conf_normalized = data.is_confidence_normalized()

    m = Model('Stealth actors')
    w_unary = 1.  # positive unary is a bonus
    pose_not_present_cost = w_unary * -1000  # negative unary is a penalty
    problem = ActorProblem2(n_actors=n_actors,
                            pose_not_present_unary_cost=pose_not_present_cost,
                            # positive pairwise is a penalty
                            pose_not_present_pw_cost=1000. * w_unary)

    objective = None
    prev_pose_in_2d = None
    prev_frame_id = None
    for frame_id in data.get_frames():
        # try:
        #     frame_id = int(frame_str.split('_')[1])
        # except ValueError:
        #     print("skipping key %s" % frame_id)
        #     continue
        frame_str = "color_%05d" % frame_id
        # if frame_id > 30:
        #     break
        # pose_in = np.array(data[frame_str][u'centered_3d'])
        # pose_in_2d = np.array(data[frame_str][u'pose_2d'])
        pose_in_2d = data.get_poses_2d(frame_id=frame_id)
        # visible = np.array(data[frame_str][u'visible'])
        # vis_f = np.array(data[frame_str][u'visible_float'])
        vis_f = data.get_confidences(frame_id=frame_id)
        assert pose_in_2d.ndim == 3, "no: %s" % repr(pose_in_2d.shape)

        problem.add_frame(frame_id,
                          n_vars=pose_in_2d.shape[0])

        # unary
        min_count = 0
        for pose_id in range(pose_in_2d.shape[0]):
            conf = vis_f[pose_id, ...]
            if not is_conf_normalized:
                conf = get_conf_thresholded(conf, thresh_log_conf=None,
                                            dtype_np=np.float32)
            cnt = np.sum(conf > 0.5)
            if cnt > conf.shape[0] // 2:
                min_count += 1
            unary = w_unary * np.sum(conf) / conf.shape[0]
            if frame_id == 251:
                print("here")
            # print("[%s] unary: %s" % (frame_id, unary))
            problem.add_unary(frame_id, pose_id, cost=unary)
        problem.set_min_count(frame_id, min_count)

        # pairwise
        if prev_pose_in_2d is not None:
            for prev_pose_id in range(prev_pose_in_2d.shape[0]):
                prev_pose = prev_pose_in_2d[prev_pose_id, :, :]
                for pose_id in range(pose_in_2d.shape[0]):
                    pose = pose_in_2d[pose_id, :, :]
                    dist = prev_pose - pose
                    dist = np.linalg.norm(dist, axis=1)
                    dist *= prev_vis_f[prev_pose_id, ...] * vis_f[pose_id, ...]
                    # lg.debug("dist: %s" % repr(dist.shape))
                    cost = np.sum(dist, axis=0)
                    # if cost > 200:
                    #     cost = 1e4

                    # cost /= 1500
                    # lg.debug("cost: %s" % cost)
                    problem.add_pw_cost(prev_frame_id, prev_pose_id,
                                        frame_id, pose_id, cost)

        prev_pose_in_2d = pose_in_2d
        prev_vis_f = vis_f
        prev_frame_id = frame_id

    gb_vars = m.addVars(problem.get_n_vars(), vtype=GRB.BINARY)
    # for lin_id in range(problem.get_n_vars()):
    #     gb_vars[lin_id].set(problem.get_init_for_lin_id(lin_id))

    # unary: we want to maximize confidence
    for lin_id, cost in problem._unary.items():
        objective -= gb_vars[lin_id] * cost

    # pairwise
    for (lin_id0, lin_id1), cost in problem._pw.items():
        objective += gb_vars[lin_id0] * gb_vars[lin_id1] * cost
    # print("NO PAIRWISE!!!")

    # a pose can only be labelled once per frame, either
    # actor0, or actor1, etc.
    for (frame_id, pose_id), lin_ids in problem._constr_p.items():
        constr = None
        for lin_id in lin_ids:
            if constr is None:
                constr = gb_vars[lin_id]
            else:
                constr += gb_vars[lin_id]
        # lg.debug("[%d] pose %d can only be %s" % (frame_id, pose_id, lin_ids))
        m.addConstr(constr <= 1)

    # an actor can only be used once per frame, either
    # pose0, or pose1, etc. Note: last actor can be used multiple times,
    # it's the "pose not present" label.
    for (frame_id, actor_id), lin_ids in problem._constr_a.items():
        constr = None
        for lin_id in lin_ids:
            if constr is None:
                constr = gb_vars[lin_id]
            else:
                constr += gb_vars[lin_id]
        # lg.debug("[%d] actor %d can only be %s"
        #          % (frame_id, actor_id, lin_ids))
        m.addConstr(constr == 1)

    # maximum number of poses chosen to be visible <= n_actors
    # for frame_id, lin_ids in problem._constr_f.items():
    #     constr = None
    #     for lin_id in lin_ids:
    #         if constr is None:
    #             constr = gb_vars[lin_id]
    #         else:
    #             constr += gb_vars[lin_id]
    #     m.addConstr(constr <= problem._n_actors)

    first_constrained = False  # type: bool
    min_frame_id = min(data.get_frames())  # type: int
    assert isinstance(min_frame_id, int)
    # anchor first pose as first actor
    if constraints and 'labels' in constraints:
        for frame_str, labels in constraints['labels'].items():
            frame_id = int(frame_str.split('_')[1])
            if isinstance(labels, list):
                assert len(labels) == n_actors, \
                    "frame: %d, %s %s" % (frame_id, len(labels), n_actors)
                # assert len(set(labels)) == len(labels), \
                #     "%s: %s" % (set(labels), labels)
                labels = {i: v for i, v in enumerate(labels)}
            for actor_id, pose_id in labels.items():
                pose_id = int(pose_id)
                if pose_id < 0:
                    pose_id = problem._max_pose_ids[frame_id]
                lin_id = problem.get_lin_id(frame_id, pose_id, actor_id)
                m.addConstr(gb_vars[lin_id] == 1)
                if not first_constrained and frame_id == min_frame_id \
                  and actor_id == 0:
                    first_constrained = True
    if not first_constrained:
        m.addConstr(gb_vars[0] == 1)
    # m.addConstr(gb_vars[36] == 1)
    # m.addConstr(gb_vars[40] == 1)

    m.setObjective(objective, GRB.MINIMIZE)
    m.Params.timeLimit = 300 if not first_run else 10
    # m.solver.callSolver(m)
    m.optimize()

    pose_ids = defaultdict(dict)
    prev_frame_id = None
    prev_lin_ids = {}
    curr_lin_ids = {}
    labelings = defaultdict(dict)
    for lin_id, v in enumerate(m.getVars()):
        frame_id, pose_id, actor_id = \
            problem.get_frame_id_pose_id_actor_id(lin_id)
        # print("[%d] %s: %s; pose %d is %sactor %s"
        #       % (frame_id, v.varName, v.x, pose_id,
        #          "not " if v.x < 0.5 else "", actor_id))
        problem._solution[lin_id] = v.x

        if prev_frame_id is not None:
            if prev_frame_id != frame_id:
                prev_lin_ids = copy.deepcopy(curr_lin_ids)
                curr_lin_ids.clear()

        # print("[#{f:d}][{l:d}] unary for p{p0:d}, a{a0:d} is {"
        #          "cost:f}{chosen:s}".format(
        #   f=frame_id, p0=pose_id, a0=actor_id,
        #   cost=problem._unary[lin_id], l=lin_id,
        #   chosen=" <-- chosen" if v.x > 0.5 else ""
        # ))

        if v.x > 0.5:
            curr_lin_ids[lin_id] = {"frame_id": frame_id,
                                    "pose_id": pose_id,
                                    "actor_id": actor_id}
            if pose_id == problem._max_pose_ids[frame_id]:
                pose_id = -1
            if frame_id in pose_ids and actor_id != n_actors:
                assert actor_id not in pose_ids[frame_id], "no"
            try:
                pose_ids[frame_id][actor_id] = pose_id
            except KeyError:
                pose_ids[frame_id] = {actor_id: pose_id}
            labelings[frame_id][pose_id] = actor_id

            # print("pw: %s" % problem._pw[lin_id])
            # for lin_id0, entries0 in prev_lin_ids.items():
            #     if (lin_id0, lin_id) in problem._pw:
            #         print("[#{f:d}] pw {l0:d}(p{p0:d},a{a0:d})"
            #               "->{l1:d}(p{p1:d},a{a1:d}) is {cost:f}".format(
            #           l0=lin_id0, l1=lin_id,
            #           cost=problem._pw[(lin_id0, lin_id)],
            #           a0=entries0['actor_id'], a1=actor_id,
            #           f=frame_id, p0=entries0['pose_id'], p1=pose_id
            #         ))


        prev_frame_id = frame_id

    # enforce constraints
    # if constraints and 'labels' in constraints:
    #     for frame_str, labels in constraints['labels'].items():
    #         frame_id = int(frame_str.split('_')[1])
    #         if isinstance(labels, list):
    #             labels = {v: i for i, v in enumerate(labels)}
    #         for pose_id, actor_id in labels.items():
    #             pose_ids[frame_id][actor_id] = int(pose_id)

    try:
        for frame_id in labelings:
            if frame_id % 5:
                continue
            print("\"color_%05d\": {%s},"
                  % (frame_id,
                     ", ".join(["\"%s\": %s" % (key, val)
                               for key, val in labelings[frame_id].items()])))
    except TypeError:
        pass


    # if we have more, pick the first...
    # if len(pose_in.shape) > 2:
    #     pose_in = pose_in[0, :, :]
    #     pose_in_2d = pose_in_2d[0, :, :]
    #     visible = visible[0]

    return pose_ids, problem, data
Esempio n. 2
0
def identify_actors(data):
    m = Model('Stealth actors')

    problem = ActorProblem()

    objective = None
    prev_pose_in_2d = None
    prev_frame_id = None
    for frame_str in sorted(data):
        try:
            frame_id = int(frame_str.split('_')[1])
        except ValueError:
            print("skipping key %s" % frame_id)
            continue

        pose_in = np.array(data[frame_str][u'centered_3d'])
        pose_in_2d = np.array(data[frame_str][u'pose_2d'])
        visible = np.array(data[frame_str][u'visible'])
        assert pose_in_2d.ndim == 3, "no: %s" % repr(pose_in_2d.shape)

        problem.add_frame(frame_id, pose_in_2d.shape[0])

        if prev_pose_in_2d is not None:
            for prev_pose_id in range(prev_pose_in_2d.shape[0]):
                prev_pose = prev_pose_in_2d[prev_pose_id, :, :]
                for pose_id in range(pose_in_2d.shape[0]):
                    pose = pose_in_2d[pose_id, :, :]
                    dist = prev_pose - pose
                    lg.debug("dist: %s" % repr(dist.shape))
                    cost = np.sum(np.linalg.norm(dist, axis=1), axis=0)
                    lg.debug("cost: %s" % cost)
                    problem.add_cost(prev_frame_id, prev_pose_id, frame_id, pose_id, cost)

        prev_pose_in_2d = pose_in_2d
        prev_frame_id = frame_id

    gb_vars = m.addVars(problem.get_n_vars(), vtype=GRB.BINARY)

    for (lin_id0, lin_id1), cost in problem._pw.items():
        # lin_id0 = problem.get_lin_id(prev_frame_id, prev_pose_id)
        # lin_id1 = problem.get_lin_id(frame_id, pose_id)
        objective += gb_vars[lin_id0] * gb_vars[lin_id1] * cost

    for frame_id, lin_ids in problem._constr.items():
        constr = None
        for lin_id in lin_ids:
            if constr is None:
                constr = gb_vars[lin_id]
            else:
                constr += gb_vars[lin_id]
        m.addConstr(constr == 1)

    m.setObjective(objective, GRB.MINIMIZE)
    # m.solver.callSolver(m)
    m.optimize()

    pose_ids = dict()
    for lin_id, v in enumerate(m.getVars()):
        print(v.varName, v.x)
        if v.x > 0.5:
            frame_id, pose_id = problem.get_frame_id_pose_id(lin_id)
            assert frame_id not in pose_ids, "no"
            pose_ids[frame_id] = pose_id

    # if we have more, pick the first...
    # if len(pose_in.shape) > 2:
    #     pose_in = pose_in[0, :, :]
    #     pose_in_2d = pose_in_2d[0, :, :]
    #     visible = visible[0]

    return pose_ids