基于深度学习的多目标跟踪算法（上）：端到端的数据关联

self.augment = Compose([
                ConvertFromInts(),
                PhotometricDistort(),
                Expand(self.mean),
                RandomSampleCrop(),
                RandomMirror(),
                ToPercentCoords(),
                Resize(self.size),
                SubtractMeans(self.mean),
                ResizeShuffleBoxes(),
                FormatBoxes(),
                ToTensor()
            ])

def forward_feature_extracter(self, x, l):



    
    '''
    extract features from the vgg layers and extra net
    :param x:
    :param l:
    :return: the features
    '''
    s = list()

    x = self.forward_vgg(x, self.vgg, s)
    x = self.forward_extras(x, self.extras, s)
    x = self.forward_selector_stacker1(s, l, self.selector)

    return x
def forward_vgg(self, x, vgg, sources):
    for k in range(16):
        x = vgg[k](x)
    sources.append(x)

    for k in range(16, 23):
        x = vgg[k](x)
    sources.append(x)

    for k in range(23, 35):
        x = vgg[k](x)
    sources.append(x)
    return x

def forward_extras(self, x, extras, sources):
    for k, v in enumerate(extras):
        x = v(x) #x = F.relu(v(x), inplace=True)        #done: relu is unnecessary.
        if k % 6 == 3:                  #done: should select the output of BatchNormalization (-> k%6==2)
            sources.append(x)
    return x

def forward_selector_stacker1(self, sources, labels, selector):
    '''
    :param sources: [B, C, H, W]
    :param labels: [B, N, 1, 1, 2]
    :return: the connected feature
    '''
    sources = [
        F.relu(net(x), inplace=True) for net, x in zip(selector, sources)
    ]

    res = list()
    for label_index in range(labels.size(1)):
        label_res = list()
        for source_index in range(len(sources)):
            # [N, B, C, 1, 1]
            label_res.append(
                # [B, C, 1, 1]
                F.grid_sample(sources[source_index],  # [B, C, H, W]
                              labels[:, label_index, :]  # [B, 1, 1, 2
                              ).squeeze(2).squeeze(2)
            )
        res.append(torch.cat(label_res, 1))

    return torch.stack(res, 1)

def forward_stacker2(self, stacker1_pre_output, stacker1_next_output):
    stacker1_pre_output = stacker1_pre_output.unsqueeze(2).repeat(1, 1, self.max_object, 1).permute(0, 3, 1, 2)
    stacker1_next_output = stacker1_next_output.unsqueeze(1).repeat(1, self.max_object, 1, 1).permute(0, 3, 1, 2)

    stacker1_pre_output = self.stacker2_bn(stacker1_pre_output.contiguous())
    stacker1_next_output = self.stacker2_bn(stacker1_next_output.contiguous())

    output = torch.cat(
        [stacker1_pre_output, stacker1_next_output],
        1
    )

    return output

def get_similarity_uv(self, t, frame_index):
    res_similarity = []
    res_uv = []
    for i, f in enumerate(t.f):
        if len(t.f) == TrackerConfig.max_track_node and i == 0:
            continue

        all_iou = self.recorder.all_iou[frame_index][f]
        all_similarity = self.recorder.all_similarity[frame_index][f]
        selected_box_index = t.uv[i, i]
        if selected_box_index == -1: # cannot find box in f frame.
            res_similarity += [0]
            res_uv += [-1]
            continue

        # combine the similarity with the iou
        selected_similarity = np.copy(all_similarity[selected_box_index, :])
        delta_f = frame_index - f
        if delta_f in TrackerConfig.min_iou_frame_gap:
            iou_index = TrackerConfig.min_iou_frame_gap.index(delta_f)
            selected_iou = (all_iou[selected_box_index, :] >= TrackerConfig.min_iou[iou_index]).astype(float)
            selected_iou = np.append(selected_iou, 1.0)
            selected_similarity = selected_similarity * selected_iou

        max_index = np.argmax(selected_similarity)
        max_value = all_similarity[selected_box_index, max_index]

        if max_index == all_similarity.shape[1] - 1: # new node
            max_index = -1
        res_uv += [int(max_index)]
        res_similarity += [float(max_value)]

    # get the representation box of this frame.
    res = {}
    for uv, s in zip(res_uv, res_similarity):
        # if s 



    
        #     continue
        if uv not in res:
            res[uv] = [s]
        else:
            res[uv] += [s]

    if len(res.keys()) > 0:
        max_uv = max(res.keys(), key=(lambda k: np.sum(res[k])))
    else:
        max_uv = -1

    res_similarity += [1]
    res_uv += [max_uv]

    if max_uv == -1:
        t.age += 1
    else:
        t.age = 0

    return res_similarity, res_uv

self.lstm_row = nn.GRU(element_dim, hidden_dim, bidirectional=biDirenction, num_layers=2)
self.lstm_col = nn.GRU(512


    
, hidden_dim, bidirectional=biDirenction, num_layers=2)

# The linear layer that maps from hidden state space to tag space
if biDirenction:
    # *2 directions * 2 ways concat
    self.hidden2tag_1 = nn.Linear(hidden_dim * 2, 256)
    self.hidden2tag_2 = nn.Linear(256, 64)
    self.hidden2tag_3 = nn.Linear(64, target_size)
else:
    # * 2 ways concat
    self.hidden2tag_1 = nn.Linear(hidden_dim, target_size)