def transform_annotation(x, orig_dim, model_dim):
"""Convert the annotation/target boxes from VoTT to a format understood by
dataset class
Arguments
---------
x : string
Consists of:
label
x_center measured from left (0-1)
y_center measured from top (0-1)
width (0-1)
height (0-1)
"""
if not x:
return None
boxes = np.array([a.rstrip().split(' ') for a in x], dtype='float32')
# Get the bounding boxes and convert them into proper format
category_ids = boxes[:, 0]
boxes = boxes[:, 1:]
boxes = np.array(boxes)
# Change y_center to measure from bottom
boxes[:, 1] = 1. - boxes[:, 1]
boxes = boxes.reshape(-1, 4)
# Center to corner
boxes = center_to_corner_2d(boxes)
boxes[:, [0, 2]] *= orig_dim[1]
boxes[:, [1, 3]] *= orig_dim[0]
category_ids = np.array(category_ids).reshape(-1, 1)
ground_truth = np.concatenate([boxes, category_ids], 1).reshape(-1, 5)
return ground_truth
def write_results(prediction,
confidence,
num_classes,
model_dim,
orig_dim,
nms=True,
nms_conf=0.7):
"""
Arguments
---------
prediction : tensor (3D)
[batch, image_id, [x_center, y_center, width, height, objectness_score, class_score1, class_score2, ...]]
Returns
--------
output : tensor (2D)
[image_id, [batch_index, x_1, y_1, x_2, y_2, objectness_score, class_index, class_probability]]
"""
# Initialize to no output
output = -1
# Technically, this should always be 1
batch_size = prediction.size(0)
# Get rid of 1st dim
orig_dim = orig_dim.squeeze(0)
# If the entire batch contains 0 for objectness score, skip
try:
torch.nonzero(prediction[:, :, 4]).transpose(0, 1).contiguous()
except:
return -1
# Keep track of if output has been compiled yet (for concatenation)
write = False
for ind in range(batch_size):
pred = prediction[ind]
if pred.shape[0] > 0:
# Get x1y1x2y2
pred = center_to_corner_2d(pred)
# #Get the class having maximum score, and the index of that class
# #Get rid of num_classes softmax scores
# #Add the class index and the class score of class having maximum score
max_conf_score, max_conf = torch.max(pred[:, 5:5 + num_classes], 1)
max_conf = max_conf.float().unsqueeze(1)
max_conf_score = max_conf_score.float().unsqueeze(1)
seq = (pred[:, :5], max_conf, max_conf_score)
image_pred = torch.cat(seq, 1)
#Get rid of the zero entries for objectness
non_zero_ind = (torch.nonzero(image_pred[:, 4]))
image_pred_ = image_pred[non_zero_ind.squeeze(), :].view(-1, 7)
# Remove low confidence by class probs
image_pred_ = image_pred_[image_pred_[:, -1] > confidence, :]
#Get the various classes detected in the image
try:
img_classes = unique(image_pred_[:, -2].int())
print('img_classes ', img_classes)
except:
continue
#WE will do NMS classwise
for label in img_classes:
#get the detections with one particular class
cls_mask_ind = (image_pred_[:, -2].int() == label)
# class_mask_ind = torch.nonzero(cls_mask[:,-2]).squeeze()
image_pred_class = image_pred_[cls_mask_ind].view(-1, 7)
#sort the detections such that the entry with the maximum objectness
#confidence is at the top
conf_sort_index = torch.sort(image_pred_class[:, 4],
descending=True)[1]
image_pred_class = image_pred_class[conf_sort_index]
idx = image_pred_class.size(0)
#if nms has to be done
if nms:
#For each detection
for i in range(idx):
#Get the IOUs of all boxes that come after the one we are looking at
#in the loop
try:
ious = bbox_iou(image_pred_class[i].unsqueeze(0),
image_pred_class[i + 1:])
except ValueError:
continue
except IndexError:
continue
# Zero out all the detections that have IoU > nms treshhold
iou_mask = (ious < nms_conf).float().unsqueeze(1)
image_pred_class[i + 1:] *= iou_mask
# Keep the non-zero entries for objectness
non_zero_ind = torch.nonzero(
image_pred_class[:, 4]).squeeze()
image_pred_class = image_pred_class[non_zero_ind].view(
-1, 7)
batch_ind = image_pred_class.new(image_pred_class.size(0),
1).fill_(ind)
seq = batch_ind, image_pred_class
if not write:
output = torch.cat(seq, 1)
write = True
else:
out = torch.cat(seq, 1)
output = torch.cat((output, out))
return output
orig_w, orig_h = img_.shape[0], img_.shape[1]
print(img_file)
print(i)
# Read image and prepare for input to network
img, orig_im, orig_dim = prep_image(plt.imread(img_file), model_dim)
orig_dim = torch.FloatTensor(orig_dim).repeat(1, 2)
# Read ground truth labels
ground_truths_file = img_file.replace(img_file.split('.')[-1], 'txt')
with open(ground_truths_file, 'r') as f:
# Read ground truth file and transform annotation to: xc,yc,w,h,label
ground_truths = transform_annotation(f.readlines(), img_.shape,
model_dim)
# Convert xc,yc,w,h --> x1,y1,x2,y2
ground_truths = center_to_corner_2d(ground_truths)
# Scale up from model size to original image size
# Go from square (0-1)*model_dim to original image shape scale
ground_truths[:, [0, 2]] *= orig_w / model_dim
ground_truths[:, [1, 3]] *= orig_h / model_dim
class_labels = ground_truths[:, -1]
num_gts += ground_truths.shape[0]
img = img.to(device)
with torch.no_grad():
output = model(img)
# NB, output is:
# [batch, image_id, [x_center, y_center, width, height, objectness_score, class_score1, class_score2, ...]]
os.makedirs(eval_output_dir, exist_ok=True)
for i, (image, ground_truth, filepath) in enumerate(test_loader):
img_file = filepath[0].rstrip()
img_ = plt.imread(img_file)
orig_h, orig_w = img_.shape[0], img_.shape[1]
orig_dim = torch.FloatTensor([orig_w, orig_h]).repeat(1,2)
print(img_file)
print(i)
# Deal with ground truth: read, convert to corners and scale
ground_truths = []
gt_file = '.'.join(img_file.split('.')[:-1]) + '.txt'
gt_df = np.array(pd.read_csv(gt_file, sep=' ', header=None))
gt_df = center_to_corner_2d(gt_df[:, 1:])
orig_dim_np = np.array([orig_w, orig_h, orig_w, orig_h])
gt_df *= orig_dim_np
ground_truths.append(gt_df)
num_gts = gt_df.shape[0]
# ground_truths = []
# if len(ground_truth) == 0:
# continue
# else:
# ground_truths.append(ground_truth)
# Predict on input test image
image = image.to(device)
with torch.no_grad():
output = model(image)