def draw_y_on_x(x, y):
for batch in range(x.shape[0]):
boxes = []
classes = []
for scale in range(len(y)):
# print(y[scale][batch, ...].shape)
for a in range(y[scale][batch, ...].shape[0]):
for w in range(y[scale][batch, ...].shape[1]):
for h in range(y[scale][batch, ...].shape[2]):
cell = y[scale][batch, a, w, h, ...]
if cell[0] == 1:
bbox = cell[1:5]
cell_number = y[scale][batch, ...].shape[1]
box_size = 416 / cell_number
bbox_yolo = torch.tensor([box_size, box_size, box_size, box_size]) * bbox + torch.tensor([box_size * h, box_size * w, 0, 0])
bbox_corners = torch.tensor([bbox_yolo[0] - bbox_yolo[2] / 2,
bbox_yolo[1] - bbox_yolo[3] / 2,
bbox_yolo[0] + bbox_yolo[2] / 2,
bbox_yolo[1] + bbox_yolo[3] / 2])
boxes.append(bbox_corners.tolist())
classes.append(int(cell[5]))
boxes = torch.tensor(boxes)
labels = [config.PASCAL_CLASSES[x] for x in classes]
x[batch, ...] = draw_bounding_boxes(image=x[batch, ...].type(torch.uint8), boxes=boxes, colors=(255, 255, 255), labels=labels, width=2)
return x
def draw_yp_on_x(x, yp, probability_threshold, anchors):
batch, channel, width, height = x.shape
for b in range(batch):
boxes = []
for scale in range(len(yp)):
yp_anchor, yp_w, yp_h, _ = yp[scale][b, ...].shape
num_cells = yp[scale][b, ...].shape[1] # 13, 26, 52
box_size = width / num_cells # width = 416
for a in range(yp_anchor):
for w in range(yp_w):
for h in range(yp_h):
# Current cell probability is above threshold
if yp[scale][b, a, w, h, 0] > probability_threshold:
# p, x, y, w, h, logit
cell = yp[scale][b, a, w, h, ...]
# p, x, y, w, h
cell_class = torch.argmax(cell[5:])
# p, x, y, w, h, c
cell_box = torch.cat((cell[0:5], torch.tensor([cell_class])), dim=0)
# Decode yp
bbox_yolo = torch.tensor([cell_box[0],
torch.sigmoid(cell_box[1]) + (h + 0.5) * box_size,
torch.sigmoid(cell_box[2]) + (w + 0.5) * box_size,
anchors[scale][a][0] * math.pow(math.e, cell_box[3]) * width,
anchors[scale][a][1] * math.pow(math.e, cell_box[4]) * height,
cell_box[5]])
# Convert bbox midpoints to corners
bbox_corner = torch.tensor([bbox_yolo[0],
bbox_yolo[1] - bbox_yolo[3] / 2,
bbox_yolo[2] - bbox_yolo[4] / 2,
bbox_yolo[1] + bbox_yolo[3] / 2,
bbox_yolo[2] + bbox_yolo[4] / 2,
bbox_yolo[5]])
boxes.append(bbox_corner)
nms_bboxes = non_maximum_suppression(bboxes=boxes, iou_threshold=0.25, threshold=0.5, box_format="corners")
boxes = [box[1:5].tolist() for box in nms_bboxes]
boxes = torch.tensor(boxes)
labels = [config.PASCAL_CLASSES[int(x[5].item())] for x in nms_bboxes]
x[b, ...] = draw_bounding_boxes(image=x[b, ...].type(torch.uint8), boxes=boxes, colors=(255, 0, 255), labels=labels, width=2)
return x
def main():
image_size = 416
anchors = torch.tensor(config.ANCHORS[0] + config.ANCHORS[1] + config.ANCHORS[2])
x = torch.rand((3, image_size, image_size))
# boxes = 416 * torch.rand((12, 4))
boxes = anchors_to_corners(anchors, image_size)
labels = ["1", "2", "3", "4", "5", "6", "7", "8", "9"]
image = draw_bounding_boxes(image=x.type(torch.uint8), boxes=boxes, colors=(255, 255, 255), labels=labels)
image = image.permute(1, 2, 0)
plt.imshow(image)
plt.show()
def main():
# Data loading
test_csv_path = config.DATASET + "100examples.csv"
test_dataset = YoloDataset(test_csv_path,
transforms=config.max_transforms,
Scale=config.Scale,
image_dir=config.IMG_DIR,
label_dir=config.LABEL_DIR,
anchor_boxes=config.ANCHORS,
number_of_anchors=3,
number_of_scales=3,
ignore_iou_threshold=0.5,
num_anchors_per_scale=3)
test_loader = DataLoader(dataset=test_dataset,
batch_size=config.BATCH_SIZE,
num_workers=config.NUM_WORKERS,
pin_memory=config.PIN_MEMORY,
shuffle=False,
drop_last=True)
# Model
model = YoloV3(num_classes=config.C)
optimizer = optim.Adam(model.parameters(),
config.LEARNING_RATE,
weight_decay=config.WEIGHT_DECAY)
loss_function = YoloLoss()
scalar = torch.cuda.amp.GradScaler()
writer = SummaryWriter()
model.to(device=config.DEVICE)
model.eval()
def load_checkpoint(checkpoint_file, model, optimizer, lr):
print("=> Loading checkpoint")
checkpoint = torch.load(checkpoint_file, map_location=config.DEVICE)
model.load_state_dict(checkpoint["state_dict"])
optimizer.load_state_dict(checkpoint["optimizer"])
# If we don't do this then it will just have learning rate of old checkpoint
# and it will lead to many hours of debugging \:
for param_group in optimizer.param_groups:
param_group["lr"] = lr
load_checkpoint(config.CHECKPOINT_FILE, model, optimizer,
config.LEARNING_RATE)
# all_pred_boxes, all_true_boxes = get_evaluation_bboxes(test_loader, model, iou_threshold=0.5, anchors=config.ANCHORS, threshold=0.5, box_format="midpoint", device="cuda")
model.eval()
train_idx = 0
anchors = config.ANCHORS
threshold = 0.9
from utils_depreciated import cells_to_bboxes
from utils_depreciated import non_max_suppression
for batch_idx, (x, labels) in enumerate(test_loader):
x = x.float()
x = x.to("cuda")
all_pred_boxes = []
all_true_boxes = []
with torch.no_grad():
predictions = model(x)
batch_size = x.shape[0]
bboxes = [[] for _ in range(batch_size)]
for i in range(3):
S = predictions[i].shape[2]
anchor = torch.tensor([*anchors[i]]).to("cuda") * S
boxes_scale_i = cells_to_bboxes(predictions[i],
anchor,
S=S,
is_preds=True)
for idx, (box) in enumerate(boxes_scale_i):
bboxes[idx] += box
# we just want one bbox for each label, not one for each scale
true_bboxes = cells_to_bboxes(labels[2], anchor, S=S, is_preds=False)
for idx in range(batch_size):
nms_boxes = non_max_suppression(bboxes[idx],
iou_threshold=0.9,
threshold=0.9,
box_format="midpoint")
for nms_box in nms_boxes:
all_pred_boxes.append([train_idx] + nms_box)
for box in true_bboxes[idx]:
if box[1] > threshold:
all_true_boxes.append([train_idx] + box)
train_idx += 1
from torchvision_utils import draw_bounding_boxes
x = x.to('cpu') * 255
#boxes = all_pred_boxes #cells_to_boxes_render(y)
boxes = [(torch.tensor(x)[3:] *
torch.tensor([416, 416, 416, 416])).tolist()
for x in all_pred_boxes]
boxes = torch.tensor(boxes)
for idx in range(x.shape[0]):
x[idx, ...] = draw_bounding_boxes(image=x[idx,
...].type(torch.uint8),
boxes=boxes,
colors=None,
width=2,
pad_value=2)
grid = torchvision.utils.make_grid(x)
writer.add_image("images", grid, global_step=batch_idx)
writer.flush()
"""