def run_demo(cfg,
ckpt,
score_threshold,
images_dir: pathlib.Path,
output_dir: pathlib.Path,
dataset_type,
num_images=None):
if dataset_type == "voc":
class_names = VOCDataset.class_names
elif dataset_type == 'coco':
class_names = COCODataset.class_names
elif dataset_type == "mnist":
class_names = MNISTDetection.class_names
elif dataset_type == "tdt4265":
class_names = TDT4265Dataset.class_names
elif dataset_type == "waymo":
class_names = WaymoDataset.class_names
else:
raise NotImplementedError('Not implemented now.')
model = SSDDetector(cfg)
model = torch_utils.to_cuda(model)
checkpointer = CheckPointer(model, save_dir=cfg.OUTPUT_DIR)
checkpointer.load(ckpt, use_latest=ckpt is None)
weight_file = ckpt if ckpt else checkpointer.get_checkpoint_file()
print('Loaded weights from {}'.format(weight_file))
image_paths = list(images_dir.glob("*.png")) + list(
images_dir.glob("*.jpg"))
output_dir.mkdir(exist_ok=True, parents=True)
transforms = build_transforms(cfg, is_train=False)
model.eval()
drawn_images = []
for i, image_path in enumerate(
tqdm.tqdm(image_paths[:num_images], desc="Predicting on images")):
image_name = image_path.stem
image = np.array(Image.open(image_path).convert("RGB"))
height, width = image.shape[:2]
images = transforms(image)[0].unsqueeze(0)
result = model(torch_utils.to_cuda(images))[0]
result = result.resize((width, height)).cpu().numpy()
boxes, labels, scores = result['boxes'], result['labels'], result[
'scores']
indices = scores > score_threshold
boxes = boxes[indices]
labels = labels[indices]
scores = scores[indices]
drawn_image = draw_boxes(image, boxes, labels, scores,
class_names).astype(np.uint8)
drawn_images.append(drawn_image)
im = Image.fromarray(drawn_image)
output_path = output_dir.joinpath(f"{image_name}.png")
im.save(output_path)
return drawn_images
def run_demo(cfg, ckpt, score_threshold, images_dir: pathlib.Path,
output_dir: pathlib.Path, dataset_type):
if dataset_type == "voc":
class_names = VOCDataset.class_names
elif dataset_type == 'coco':
class_names = COCODataset.class_names
elif dataset_type == "mnist":
class_names = MNISTDetection.class_names
else:
raise NotImplementedError('Not implemented now.')
model = SSDDetector(cfg)
model = torch_utils.to_cuda(model)
checkpointer = CheckPointer(model, save_dir=cfg.OUTPUT_DIR)
checkpointer.load(ckpt, use_latest=ckpt is None)
weight_file = ckpt if ckpt else checkpointer.get_checkpoint_file()
print('Loaded weights from {}'.format(weight_file))
image_paths = list(images_dir.glob("*.png")) + list(
images_dir.glob("*.jpg"))
output_dir.mkdir(exist_ok=True, parents=True)
transforms = build_transforms(cfg, is_train=False)
model.eval()
drawn_images = []
for i, image_path in enumerate(image_paths):
start = time.time()
image_name = image_path.name
image = np.array(Image.open(image_path).convert("RGB"))
height, width = image.shape[:2]
images = transforms(image)[0].unsqueeze(0)
load_time = time.time() - start
start = time.time()
result = model(torch_utils.to_cuda(images))[0]
inference_time = time.time() - start
result = result.resize((width, height)).cpu().numpy()
boxes, labels, scores = result['boxes'], result['labels'], result[
'scores']
indices = scores > score_threshold
boxes = boxes[indices]
labels = labels[indices]
scores = scores[indices]
meters = "|".join([
'objects {:02d}'.format(len(boxes)),
'load {:03d}ms'.format(round(load_time * 1000)),
'inference {:03d}ms'.format(round(inference_time * 1000)),
'FPS {}'.format(round(1.0 / inference_time))
])
image_name = image_path.name
drawn_image = draw_boxes(image, boxes, labels, scores,
class_names).astype(np.uint8)
drawn_images.append(drawn_image)
return drawn_images
def get_detections(cfg, ckpt):
model = SSDDetector(cfg)
model = torch_utils.to_cuda(model)
checkpointer = CheckPointer(cfg, model, save_dir=cfg.OUTPUT_DIR)
checkpointer.load(ckpt, use_latest=ckpt is None)
weight_file = ckpt if ckpt else checkpointer.get_checkpoint_file()
print('Loaded weights from {}'.format(weight_file))
dataset_path = DatasetCatalog.DATASETS["tdt4265_test"]["data_dir"]
dataset_path = pathlib.Path(cfg.DATASET_DIR, dataset_path)
image_dir = pathlib.Path(dataset_path, "images")
image_paths = list(image_dir.glob("*.jpg"))
transforms = build_transforms(cfg, is_train=False)
model.eval()
detections = []
labels = read_labels(
image_dir.parent.parent.joinpath("train", "labels.json"))
check_all_images_exists(labels, image_paths)
# Filter labels on if they are test and only take the 7th frame
labels = [label for label in labels if label["is_test"]]
labels = [label for label in labels if label["image_id"] % 7 == 0]
for i, label in enumerate(tqdm.tqdm(labels, desc="Inference on images")):
image_id = label["image_id"]
image_path = image_dir.joinpath(f"{image_id}.jpg")
image_detections = {"image_id": int(image_id), "bounding_boxes": []}
image = np.array(Image.open(image_path).convert("RGB"))
height, width = image.shape[:2]
images = transforms(image)[0].unsqueeze(0)
result = model(torch_utils.to_cuda(images))[0]
result = result.resize((width, height)).cpu().numpy()
boxes, labels, scores = result['boxes'], result['labels'], result[
'scores']
for idx in range(len(boxes)):
box = boxes[idx]
label_id = labels[idx]
label = TDT4265Dataset.class_names[label_id]
assert label != "__background__"
score = float(scores[idx])
assert box.shape == (4, )
json_box = {
"xmin": float(box[0]),
"ymin": float(box[1]),
"xmax": float(box[2]),
"ymax": float(box[3]),
"label": str(label),
"label_id": int(label_id),
"confidence": float(score)
}
image_detections["bounding_boxes"].append(json_box)
detections.append(image_detections)
return detections
def get_detections(cfg, ckpt):
model = SSDDetector(cfg)
model = torch_utils.to_cuda(model)
checkpointer = CheckPointer(model, save_dir=cfg.OUTPUT_DIR)
checkpointer.load(ckpt, use_latest=ckpt is None)
weight_file = ckpt if ckpt else checkpointer.get_checkpoint_file()
print('Loaded weights from {}'.format(weight_file))
dataset_path = DatasetCatalog.DATASETS["tdt4265_test"]["data_dir"]
dataset_path = pathlib.Path(cfg.DATASET_DIR, dataset_path)
image_dir = pathlib.Path(dataset_path)
image_paths = list(image_dir.glob("*.jpg"))
transforms = build_transforms(cfg, is_train=False)
model.eval()
detections = []
for image_path in tqdm.tqdm(image_paths, desc="Inference on images"):
image = np.array(Image.open(image_path).convert("RGB"))
height, width = image.shape[:2]
images = transforms(image)[0].unsqueeze(0)
result = model(torch_utils.to_cuda(images))[0]
result = result.resize((width, height)).cpu().numpy()
boxes, labels, scores = result['boxes'], result['labels'], result[
'scores']
for idx in range(len(boxes)):
box = boxes[idx]
label_id = labels[idx]
label = TDT4265Dataset.class_names[label_id]
assert label != "__background__"
score = float(scores[idx])
assert box.shape == (4, )
xmin, ymin, xmax, ymax = box
width = xmax - xmin
height = ymax - ymin
detections.append({
"image_id": image_path.stem,
"category_id": LABEL_MAP[label],
"score": score,
"bbox": [xmin, ymin, width, height]
})
return detections
def evaluation(cfg, ckpt, N_images: int):
model = SSDDetector(cfg)
logger = logging.getLogger("SSD.inference")
checkpointer = CheckPointer(model, save_dir=cfg.OUTPUT_DIR, logger=logger)
model = torch_utils.to_cuda(model)
checkpointer.load(ckpt, use_latest=ckpt is None)
model.eval()
data_loaders_val = make_data_loader(cfg, is_train=False)
for data_loader in data_loaders_val:
batch = next(iter(data_loader))
images, targets, image_ids = batch
images = torch_utils.to_cuda(images)
imshape = list(images.shape[2:])
# warmup
print("Checking runtime for image shape:", imshape)
for i in range(10):
model(images)
start_time = time.time()
for i in range(N_images):
outputs = model(images)
total_time = time.time() - start_time
print("Runtime for image shape:", imshape)
print("Total runtime:", total_time)
print("FPS:", N_images / total_time)