示例#1
0
文件: yolo.py 项目: RI-SE/smirk
 def _forward_augment(self, x):
     img_size = x.shape[-2:]  # height, width
     s = [1, 0.83, 0.67]  # scales
     f = [None, 3, None]  # flips (2-ud, 3-lr)
     y = []  # outputs
     for si, fi in zip(s, f):
         xi = scale_img(x.flip(fi) if fi else x, si, gs=int(self.stride.max()))
         yi = self._forward_once(xi)[0]  # forward
         # cv2.imwrite(f'img_{si}.jpg', 255 * xi[0].cpu().numpy().transpose((1, 2, 0))[:, :, ::-1])  # save
         yi = self._descale_pred(yi, fi, si, img_size)
         y.append(yi)
     y = self._clip_augmented(y)  # clip augmented tails
     return torch.cat(y, 1), None  # augmented inference, train
示例#2
0
 def forward(self, x, augment=False, profile=False):
     if augment:
         img_size = x.shape[-2:]  # height, width
         s = [1, 0.83, 0.67]  # scales
         f = [None, 3, None]  # flips (2-ud, 3-lr)
         y = []  # outputs
         for si, fi in zip(s, f):
             xi = scale_img(x.flip(fi) if fi else x, si)
             yi = self.forward_once(xi)[0]  # forward
             # cv2.imwrite('img%g.jpg' % s, 255 * xi[0].numpy().transpose((1, 2, 0))[:, :, ::-1])  # save
             yi[..., :4] /= si  # de-scale
             if fi == 2:
                 yi[..., 1] = img_size[0] - yi[..., 1]  # de-flip ud
             elif fi == 3:
                 yi[..., 0] = img_size[1] - yi[..., 0]  # de-flip lr
             y.append(yi)
         return torch.cat(y, 1), None  # augmented inference, train
     else:
         return self.forward_once(x, profile)  # single-scale inference, train