| 32 | |
| 33 | |
| 34 | def predict_transform(prediction, inp_dim, anchors, num_classes, CUDA = True): |
| 35 | batch_size = prediction.size(0) |
| 36 | stride = inp_dim // prediction.size(2) |
| 37 | grid_size = inp_dim // stride |
| 38 | bbox_attrs = 5 + num_classes |
| 39 | num_anchors = len(anchors) |
| 40 | |
| 41 | anchors = [(a[0]/stride, a[1]/stride) for a in anchors] |
| 42 | |
| 43 | prediction = prediction.view(batch_size, bbox_attrs*num_anchors, grid_size*grid_size) |
| 44 | prediction = prediction.transpose(1, 2).contiguous() |
| 45 | prediction = prediction.view(batch_size, grid_size*grid_size*num_anchors, bbox_attrs) |
| 46 | |
| 47 | # Sigmoid the centre_X, centre_Y. and object confidencce |
| 48 | prediction[:, :, 0] = torch.sigmoid(prediction[:, :, 0]) |
| 49 | prediction[:, :, 1] = torch.sigmoid(prediction[:, :, 1]) |
| 50 | prediction[:, :, 4] = torch.sigmoid(prediction[:, :, 4]) |
| 51 | |
| 52 | # Add the center offsets |
| 53 | grid_len = np.arange(grid_size) |
| 54 | a, b = np.meshgrid(grid_len, grid_len) |
| 55 | |
| 56 | x_offset = torch.FloatTensor(a).view(-1, 1) |
| 57 | y_offset = torch.FloatTensor(b).view(-1, 1) |
| 58 | |
| 59 | if CUDA: |
| 60 | x_offset = x_offset.cuda() |
| 61 | y_offset = y_offset.cuda() |
| 62 | |
| 63 | x_y_offset = torch.cat((x_offset, y_offset), 1).repeat(1, num_anchors).view(-1, 2).unsqueeze(0) |
| 64 | |
| 65 | prediction[:, :, :2] += x_y_offset |
| 66 | |
| 67 | # log space transform height and the width |
| 68 | anchors = torch.FloatTensor(anchors) |
| 69 | |
| 70 | if CUDA: |
| 71 | anchors = anchors.cuda() |
| 72 | |
| 73 | anchors = anchors.repeat(grid_size*grid_size, 1).unsqueeze(0) |
| 74 | prediction[:, :, 2:4] = torch.exp(prediction[:, :, 2:4])*anchors |
| 75 | |
| 76 | # Softmax the class scores |
| 77 | prediction[:, :, 5: 5 + num_classes] = torch.sigmoid((prediction[:, :, 5: 5 + num_classes])) |
| 78 | |
| 79 | prediction[:, :, :4] *= stride |
| 80 | |
| 81 | return prediction |
| 82 | |
| 83 | |
| 84 | def load_classes(namesfile): |