def forward(self, x):
        x = x.float()
        x = cuda(x)

        x, scale = self.process_im(x)
        x = x - cuda(self.Mean)
        x = x.permute(0, 3, 1, 2)

        img_size = x.shape[2:]

        x = self.features(x)
        rpn_logits, rpn_loc = self.rpn(x)
        map_H, map_W = x.shape[2:]
        tanchors = self.anchors[:map_H, :map_W]
        tanchors = tanchors.contiguous().view(-1, 4)
        tanchors = cuda(tanchors)
        roi = self.roi_layer(rpn_loc,
                             F.softmax(rpn_logits, dim=-1)[:, 1], tanchors,
                             img_size)

        x = self.roialign(x, roi)
        fast_logits, fast_loc = self.fast(x)
        fast_loc = fast_loc * cuda(torch.tensor([0.1, 0.1, 0.2, 0.2]))
        pre = predict(
            fast_loc,
            F.softmax(fast_logits, dim=-1),
            roi[:, 1:],
            img_size[0],
            img_size[1],
        )
        pre[:, :4] = pre[:, :4] / scale
        return pre
    def forward(self, x):

        x = cuda(x.float())

        x, scale = self.process_im(x)
        x = x - cuda(self.Mean)
        x = x[..., [2, 1, 0]]
        x = x.permute(0, 3, 1, 2)

        img_size = x.shape[2:]

        C = self.features(x)
        feature_dict = OrderedDict()
        for i in range(len(C)):
            feature_dict['feat%d' % i] = C[i]

        # print('=========================')
        feature_dict = self.fpn(feature_dict)
        # for i in list(feature_dict.values()):
        #     print(i.shape)
        rpn_logits, rpn_loc = self.rpn(list(feature_dict.values()))

        tanchors = []
        map_HW = []
        for i in range(5):
            if i == 4:
                H, W = feature_dict['pool'].shape[2:4]
            else:
                H, W = feature_dict['feat%d' % i].shape[2:4]
            map_HW.append((H, W))
            tanchors.append(self.anchors[i][:H, :W].contiguous().view(-1, 4))
        tanchors = cuda(torch.cat(tanchors, dim=0))
        roi = self.roi_layer(rpn_loc.data, F.softmax(rpn_logits.data, dim=-1)[:, 1], tanchors, img_size,
                             map_HW)

        x = self.roialign_7(feature_dict, [roi], [tuple(img_size)])

        fast_logits, fast_loc = self.fast(x)
        fast_loc = fast_loc * cuda(torch.tensor([0.1, 0.1, 0.2, 0.2]))
        pre = predict(fast_loc, F.softmax(fast_logits, dim=-1), roi, img_size[0], img_size[1], iou_thresh_=0.5,
                      c_thresh=0.05)[:100]
        if pre.shape[0] == 0:
            return pre, cuda(torch.zeros((0, 28, 28)))

        roi = pre[:100]
        inds_b = roi[:, -1].long() + 1

        net_mask = self.roialign_14(feature_dict, [roi[:, :4]], [tuple(img_size)])
        net_mask = self.mask_net(net_mask)
        net_mask = torch.sigmoid(net_mask)
        inds_a = torch.arange(roi.shape[0])

        mask = net_mask[inds_a, inds_b]
        pre[:, :4] = pre[:, :4] / scale
        return pre, mask
Exemplo n.º 3
0
    def forward(self, x):
        x = x.float()
        x = cuda(x)
        x, scale = self.process_im(x)
        x = x - cuda(self.Mean)
        x = x[..., [2, 1, 0]]
        x = x.permute(0, 3, 1, 2)
        img_size = x.shape[2:]
        C = self.features(x)
        feature_dict = OrderedDict()
        for i in range(len(C)):
            feature_dict['feat%d' % i] = C[i]

        feature_dict = self.fpn(feature_dict)

        rpn_logits, rpn_loc = self.rpn(list(feature_dict.values()))
        tanchors = []
        map_HW = []
        for i in range(5):
            if i == 4:
                H, W = feature_dict['pool'].shape[2:4]
            else:
                H, W = feature_dict['feat%d' % i].shape[2:4]
            map_HW.append((H, W))
            tanchors.append(self.anchors[i][:H, :W].contiguous().view(-1, 4))
        tanchors = cuda(torch.cat(tanchors, dim=0))

        roi = self.roi_layer(rpn_loc.data,
                             F.softmax(rpn_logits.data, dim=-1)[:, 1],
                             tanchors, img_size, map_HW)

        x = self.roialign(feature_dict, [roi], [tuple(img_size)])
        fast_logits, fast_loc = self.fast(x)
        fast_loc = fast_loc * cuda(torch.tensor([0.1, 0.1, 0.2, 0.2]))
        pre = predict(
            fast_loc,
            F.softmax(fast_logits, dim=-1),
            roi,
            img_size[0],
            img_size[1],
        )
        pre[:, :4] = pre[:, :4] / scale
        return pre
    def forward(self, x):

        x = cuda(x.float())
        x, scale = self.process_im(x)
        x = x - cuda(self.Mean)
        x = x[..., [2, 1, 0]]
        x = x.permute(0, 3, 1, 2)
        img_size = x.shape[2:]

        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu(x)
        x = self.maxpool(x)
        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        rpn_logits, rpn_loc = self.rpn(x)
        map_H, map_W = x.shape[2:]
        tanchors = self.anchors[:map_H, :map_W]
        tanchors = cuda(tanchors.contiguous().view(-1, 4))

        roi = self.roi_layer(rpn_loc,
                             F.softmax(rpn_logits, dim=-1)[:, 1], tanchors,
                             img_size)

        x = self.roialign(x, roi)
        x = self.layer4(x)
        x = self.avgpool(x)
        # x = x.view(x.size(0), -1)
        fast_logits, fast_loc = self.fast(x)
        fast_loc = fast_loc * cuda(torch.tensor([0.1, 0.1, 0.2, 0.2]))
        pre = predict(
            fast_loc,
            F.softmax(fast_logits, dim=-1),
            roi[:, 1:],
            img_size[0],
            img_size[1],
        )
        pre[:, :4] = pre[:, :4] / scale
        return pre