Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
206 changes: 206 additions & 0 deletions code/libs/engine_vis.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,206 @@
import os
import json
import time

import torch
from pycocotools.coco import COCO
from pycocotools.cocoeval import COCOeval

from .utils import AverageMeter, convert_to_xywh

from torchvision.transforms.functional import to_pil_image
from PIL import ImageDraw


save_score_threshold = 0.8


cls_names = (
"aeroplane",
"bicycle",
"bird",
"boat",
"bottle",
"bus",
"car",
"cat",
"chair",
"cow",
"diningtable",
"dog",
"horse",
"motorbike",
"person",
"pottedplant",
"sheep",
"sofa",
"train",
"tvmonitor",
)

def train_one_epoch(
train_loader,
model,
optimizer,
scheduler,
curr_epoch,
device,
tb_writer=None,
print_freq=10,
):
"""Training the model for one epoch"""
# set up meters
batch_time = AverageMeter()
losses_tracker = {}
# number of iterations per epoch
num_iters = len(train_loader)
# switch to train mode
model.train()

# main training loop
print("\n[Train]: Epoch {:d} started".format(curr_epoch))
start = time.time()
for iter_idx, (imgs, targets) in enumerate(train_loader, 0):
imgs = list(img.to(device) for img in imgs)
targets = [{k: v.to(device) for k, v in t.items()} for t in targets]
# zero out optimizer
optimizer.zero_grad()
# forward / backward the model
losses = model(imgs, targets)
losses["final_loss"].backward()
# step optimizer / scheduler
optimizer.step()
scheduler.step()

# printing (only check the stats when necessary to avoid extra cost)
if (iter_idx != 0) and (iter_idx % print_freq) == 0:
# measure elapsed time (sync all kernels)
torch.cuda.synchronize()
batch_time.update((time.time() - start) / print_freq)
start = time.time()

# track all losses
for key, value in losses.items():
# init meter if necessary
if key not in losses_tracker:
losses_tracker[key] = AverageMeter()
# update
losses_tracker[key].update(value.item())

# log to tensorboard
lr = scheduler.get_last_lr()[0]
global_step = curr_epoch * num_iters + iter_idx
if tb_writer is not None:
# learning rate (after stepping)
tb_writer.add_scalar("train/learning_rate", lr, global_step)
# all losses
tag_dict = {}
for key, value in losses_tracker.items():
if key != "final_loss":
tag_dict[key] = value.val
tb_writer.add_scalars("train/all_losses", tag_dict, global_step)
# final loss
tb_writer.add_scalar(
"train/final_loss", losses_tracker["final_loss"].val, global_step
)

# print to terminal
block1 = "Epoch: [{:03d}][{:05d}/{:05d}]".format(
curr_epoch, iter_idx, num_iters
)
block2 = "Time {:.2f} ({:.2f})".format(batch_time.val, batch_time.avg)
block3 = "Loss {:.2f} ({:.2f})\n".format(
losses_tracker["final_loss"].val, losses_tracker["final_loss"].avg
)
block4 = ""
for key, value in losses_tracker.items():
if key != "final_loss":
block4 += "\t{:s} {:.2f} ({:.2f})".format(key, value.val, value.avg)

print("\t".join([block1, block2, block3, block4]))

# finish up and print
lr = scheduler.get_last_lr()[0]
print("[Train]: Epoch {:d} finished with lr={:.8f}\n".format(curr_epoch, lr))
return


def evaluate(val_loader, model, output_file, gt_json_file, device, print_freq=10):
"""Test the model on the validation set"""
# an output file will be used to save all results
assert output_file is not None

# set up meters
batch_time = AverageMeter()
# switch to evaluate mode
model.eval()
cpu_device = torch.device("cpu")

# loop over validation set
start = time.time()
det_results = []
for iter_idx, data in enumerate(val_loader, 0):
imgs, targets = data
imgs = list(img.to(device) for img in imgs)
# forward the model (wo. grad)
with torch.no_grad():
outputs = model(imgs, None)

# unpack the results
outputs = [{k: v.to(cpu_device) for k, v in t.items()} for t in outputs]
for idx, (target, output) in enumerate(zip(targets, outputs)):
image_id = target["image_id"].item()
boxes = convert_to_xywh(output["boxes"]).tolist()
scores = output["scores"].tolist()
labels = output["labels"].tolist()
pil_imgs = to_pil_image(imgs[idx])

max_score = -1

for box, score, label, outputBoxes in zip(boxes, scores, labels, output["boxes"]):
det_results.append(
{
"image_id": image_id,
"category_id": int(label),
"bbox": box,
"score": score,
}
)

if(score >= save_score_threshold):
max_score = max(score, max_score)
draw = ImageDraw.Draw(pil_imgs)
draw.rectangle(outputBoxes.numpy())
print(outputBoxes)
draw.text(outputBoxes[2:].numpy(),cls_names[int(label)-1],fill=(255,0,0))

if max_score > save_score_threshold:
pil_imgs.save('result-image-' + str(image_id) + '.jpeg')

# printing
if (iter_idx != 0) and iter_idx % (print_freq) == 0:
# measure elapsed time (sync all kernels)
torch.cuda.synchronize()
batch_time.update((time.time() - start) / print_freq)
start = time.time()

# print timing
print(
"Test: [{0:05d}/{1:05d}]\t"
"Time {batch_time.val:.2f} ({batch_time.avg:.2f})".format(
iter_idx, len(val_loader), batch_time=batch_time
)
)

# save results to json file
with open(output_file, "w") as outfile:
json.dump(det_results, outfile)

# use COCO API for evaluation
coco_gt = COCO(gt_json_file)
coco_dt = coco_gt.loadRes(output_file)
cocoEval = COCOeval(coco_gt, coco_dt, "bbox")
cocoEval.evaluate()
cocoEval.accumulate()
cocoEval.summarize()
return
28 changes: 12 additions & 16 deletions code/libs/model.py
Original file line number Diff line number Diff line change
Expand Up @@ -402,7 +402,7 @@ def compute_loss(

for image, targets_per_image in enumerate(targets):
target_boxes = targets_per_image['boxes'] # N_boxes x 4 (x1,y1,x2,y2)
target_label = targets_per_image['labels'] # N_boxes x 1 (class in range [1:20])
target_label = targets_per_image['labels'] # N_boxes x 1 (class in range [0:19])
target_areas = (target_boxes[:,2] - target_boxes[:,0]) * (target_boxes[:,3] - target_boxes[:,1]) # N_boxes x 1 (area)
target_centers = (target_boxes[:,:2] + target_boxes[:,2:]) / 2 # N_boxes x 2 (x,y)

Expand All @@ -413,25 +413,25 @@ def compute_loss(
reg_outputs_per_point = layer_reg_outputs[image] # H x W x 4
ctr_logits_per_point = layer_ctr_logits[image] # H x W x 1

W, H = layer_points.shape[:2] # layer_points - H x W x 2 (x,y)
H, W = layer_points.shape[:2] # layer_points - H x W x 2 (x,y)
N_boxes = target_boxes.shape[0] # target_boxes - N_boxes x 4 (x1,y1,x2,y2)

center_boxes_x1y1 = target_centers - self.center_sampling_radius * layer_stride # N_boxes x 2
center_boxes_x2y2 = target_centers + self.center_sampling_radius * layer_stride # N_boxes x 2
target_center_boxes = torch.concat((center_boxes_x1y1, center_boxes_x2y2), dim = 1) # N_boxes x 4

repeated_layer_points = layer_points.unsqueeze(dim=0).repeat(N_boxes, 1, 1, 1) # convert layer_points from H*W*2 to N_boxes*H*W*2
repeated_target_subboxes = target_center_boxes.view(-1, 1, 1, 4).repeat(1, W, H, 1) # convert target_center_boxes from N_boxes*4 to N_boxes*H*W*4
repeated_target_subboxes = target_center_boxes.view(-1, 1, 1, 4).repeat(1, H, W, 1) # convert target_center_boxes from N_boxes*4 to N_boxes*H*W*4

point_x = repeated_layer_points[:,:,:,0] # N_boxes x H x W
point_y = repeated_layer_points[:,:,:,1] # N_boxes x H x W
point_y = repeated_layer_points[:,:,:,0] # N_boxes x H x W
point_x = repeated_layer_points[:,:,:,1] # N_boxes x H x W

subbox_x1 = repeated_target_subboxes[:,:,:,0] # N_boxes x H x W
subbox_y1 = repeated_target_subboxes[:,:,:,1] # N_boxes x H x W
subbox_x2 = repeated_target_subboxes[:,:,:,2] # N_boxes x H x W
subbox_y2 = repeated_target_subboxes[:,:,:,3] # N_boxes x H x W

repeated_target_boxes = target_boxes.view(-1, 1, 1, 4).repeat(1, W, H, 1)
repeated_target_boxes = target_boxes.view(-1, 1, 1, 4).repeat(1, H, W, 1)
target_box_x1 = repeated_target_boxes[:,:,:,0] # N_boxes x H x W
target_box_y1 = repeated_target_boxes[:,:,:,1] # N_boxes x H x W
target_box_x2 = repeated_target_boxes[:,:,:,2] # N_boxes x H x W
Expand Down Expand Up @@ -520,16 +520,16 @@ def compute_loss(
predicted_r = reg_outputs_per_point[foreground_mask][:,2]
predicted_b = reg_outputs_per_point[foreground_mask][:,3]

foreground_x = layer_points[foreground_mask][:,0]
foreground_y = layer_points[foreground_mask][:,1]
foreground_y = layer_points[foreground_mask][:,0]
foreground_x = layer_points[foreground_mask][:,1]

predicted_x1 = foreground_x - (predicted_l) * layer_stride
predicted_y1 = foreground_y - (predicted_t) * layer_stride
predicted_x2 = foreground_x + (predicted_r) * layer_stride
predicted_y2 = foreground_y + (predicted_b) * layer_stride
predicted_xyxy = torch.stack((predicted_x1, predicted_y1, predicted_x2, predicted_y2), dim=1) # (N_foreground, 4)

target_xyxy = torch.zeros((*box_per_point.shape, 4), device=device) # (H x W x 4)
target_xyxy = torch.zeros((*box_per_point.shape, 4), device=device) # (H x W x 4)
target_xyxy[foreground_mask] = target_boxes[box_per_point[foreground_mask]] # (H x W x 4)
target_xyxy = target_xyxy[foreground_mask] # (N_foreground, 4)
target_xyxy.detach()
Expand Down Expand Up @@ -565,7 +565,7 @@ def compute_loss(

# ctr_logits_per_point : (H x W x 1)
predicted_centerness = ctr_logits_per_point[foreground_mask] # (N_foreground,1)
ctr_loss.append(binary_cross_entropy_with_logits(predicted_centerness, target_centerness))
ctr_loss.append(binary_cross_entropy_with_logits(predicted_centerness, target_centerness, reduction="sum"))


# print(cls_loss, reg_loss, ctr_loss)
Expand Down Expand Up @@ -769,19 +769,17 @@ def inference(
boxes = []
scores = []
labels = []
combination = []
for layer_stride, box_regression_per_image_per_level, logits_per_image_per_level, box_ctrness_image_per_level, points_per_level in zip(strides, reg_outputs, cls_logits, ctr_logits, points):

box_regression_per_image = box_regression_per_image_per_level[index]
logits_per_image = logits_per_image_per_level[index]
box_ctrness_per_image = box_ctrness_image_per_level[index]

scores_per_level = torch.sqrt(
torch.sigmoid(logits_per_image) * torch.sigmoid(box_ctrness_per_image)).flatten()
print(scores_per_level.shape)
torch.sigmoid(logits_per_image) * torch.sigmoid(box_ctrness_per_image)
).flatten()

keep_idxs = scores_per_level > self.score_thresh
scores_per_level = scores_per_level
topk_idxs = torch.where(keep_idxs)[0] ##INDEXES

#TODO: Decode all boxes , then filtering ( reverse this )
Expand All @@ -804,14 +802,12 @@ def inference(
labels = torch.cat(labels, dim=0)

scores, idx = scores.topk(min(scores.size(0), self.topk_candidates))
print(idx.shape, 'idx count')
boxes = boxes[idx]
labels = labels[idx]

filtered_set = batched_nms(boxes, scores, labels, self.nms_thresh)

filtered_set = filtered_set[: self.detections_per_img]
print(labels[filtered_set].shape, "final length")
detections.append(
{
'boxes': boxes[filtered_set],
Expand Down
Loading