def _test_attention_classification(model: AttentionClassificationModel, batch_size: int, n_classes=11,
image_size=(96, 96)):
model = try_cuda(model)
expected_prob_tensor_shape = (batch_size, n_classes)
expected_label_tensor_shape = (batch_size,)
expected_attention_shape_length = 3 # (batch size, height, width)
test_tensor = try_cuda(gen_random_tensor(batch_size, image_size))
labels, probs, attention_map = model.predict_label_and_heatmap(test_tensor)
assert_tensor_shape(labels, expected_label_tensor_shape, "output label shape")
assert_tensor_shape(probs, expected_prob_tensor_shape, "output prob shape")
assert len(attention_map.shape) == expected_attention_shape_length
def _test_segmentation(model: SegmentationModel,
batch_size: int,
n_classes=11,
image_size=(256, 256)):
model = try_cuda(model)
expected_prob_tensor_shape = (batch_size, n_classes) + image_size
expected_label_tensor_shape = (batch_size, ) + image_size
test_tensor = try_cuda(gen_random_tensor(batch_size, image_size))
labels, probs = model.predict_index_image(test_tensor)
assert_tensor_shape(labels, expected_label_tensor_shape,
"output label shape")
assert_tensor_shape(probs, expected_prob_tensor_shape, "output prob shape")
def test_efficientdet():
model: EfficientDetector = try_cuda(EfficientDetector(n_classes=n_classes, score_threshold=1e-2))
test_image = gen_random_tensor(1, image_size)
result = model.predict_bboxes(test_image)
assert result[0].ndim == 2 and result[0].shape[
-1] == 6, "Detection model result must contain (xmin, ymin, xmax, ymax, class, score)"
def _test_classification(model: ClassificationModel, batch_size: int, n_classes=11, image_size=(96, 96)):
expected_prob_tensor_shape = (batch_size, n_classes)
expected_label_tensor_shape = (batch_size,)
test_tensor = try_cuda(gen_random_tensor(batch_size, image_size))
labels, probs = model.predict_label(test_tensor)
assert_tensor_shape(labels, expected_label_tensor_shape, "output label shape")
assert_tensor_shape(probs, expected_prob_tensor_shape, "output prob shape")
def get_logger(log_prefix: str, args, model: BaseDeepextModel):
test_tensor = try_cuda(torch.randn(1, 3, args.image_size, args.image_size))
if args.log_type == "mlflow":
logger = MLFlowLogger(
experiment_name=
f"{log_prefix}_{args.dataset}_{model.generate_model_name()}")
# Log the model
# with mlflow.start_run():
# mlflow.pytorch.log_model(model, "model")
#
# # convert to scripted model and log the model
# scripted_pytorch_model = torch.jit.script(model)
# mlflow.pytorch.log_model(scripted_pytorch_model, "scripted_model")
return logger
elif args.log_type == "tensorboard":
logger = TensorBoardLogger(
save_dir="tensorboard_logs",
version="v",
name=
f"segmentation_demo_{args.dataset}_{model.generate_model_name()}")
logger.experiment.add_graph(model, test_tensor)
return logger
raise RuntimeError(f"Invalid log type: {args.log_type}")
Path(result_dir_path).mkdir()
label_names, label_dict = create_label_list_and_dict(label_file_path)
n_classes = len(label_names)
transforms = AlbumentationsClsWrapperTransform(
A.Compose([
A.Resize(width=width, height=height),
ToTensorV2(),
]))
dataset = ImageOnlyDataset(image_dir=images_dir_path,
image_transform=transforms)
print("Loading model...")
model: AttentionClassificationModel = try_cuda(
AttentionBranchNetwork(
n_classes=n_classes,
backbone=BackBoneKey.RESNET_18).load_from_checkpoint(checkpoint_path))
print("Model loaded")
for i, img_tensor in enumerate(tqdm.tqdm(dataset)):
origin_image = normalize255(tensor_to_cv(img_tensor))
label, prob, attention_map = model.predict_label_and_heatmap(
to_4dim(img_tensor))
attention_map = normalize255(tensor_to_cv(attention_map[0]))
blend_img = model.generate_heatmap_image(origin_image, attention_map)
result_image = cv2.resize(blend_img, dataset.current_image_size())
cv2.imwrite(f"{result_dir_path}/{label_names[label[0]]}_{i}.jpg",
result_image)
train_dataset, test_dataset = DatasetSplitter().split_train_test(test_ratio,
root_dataset,
train_transforms=train_transforms,
test_transforms=test_transforms)
else:
train_dataset = IndexImageDataset.create(image_dir_path=train_images_dir,
index_image_dir_path=train_annotations_dir,
transforms=train_transforms)
test_dataset = IndexImageDataset.create(image_dir_path=test_images_dir, index_image_dir_path=test_annotations_dir,
transforms=test_transforms)
train_data_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
test_data_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=True)
# TODO Model detail params
model: SegmentationModel = try_cuda(
ShelfNet(n_classes=n_classes, out_size=(height, width)))
if load_checkpoint_path and load_checkpoint_path != "":
model = model.load_from_checkpoint(load_checkpoint_path)
# TODO Train detail params
# Metrics/Callbacks
val_every_n_epoch = 5
callbacks = [ModelCheckpoint(period=val_every_n_epoch, filename=f"{model.generate_model_name()}",
dirpath=save_checkpoint_dir_path, monitor='val_iou', verbose=True, mode="max"),
GenerateSegmentationImageCallback(output_dir=progress_dir, per_epoch=5, model=model,
dataset=test_dataset)]
logger = MLFlowLogger(experiment_name=f"segmentation_{model.generate_model_name()}")
# Training.
Trainer(max_epochs=epoch, callbacks=callbacks, gpus=-1,
annotation_dir_path=test_annotations_dir,
class_index_dict=class_index_dict)
train_data_loader = DataLoader(train_dataset,
batch_size=batch_size,
shuffle=True,
collate_fn=AdjustDetectionTensorCollator())
test_data_loader = DataLoader(test_dataset,
batch_size=batch_size,
shuffle=True,
collate_fn=AdjustDetectionTensorCollator())
# TODO Model detail params
model: DetectionModel = try_cuda(
EfficientDetector(n_classes=n_classes,
lr=lr,
network=f"efficientdet-d0",
score_threshold=0.5))
if load_checkpoint_path and load_checkpoint_path != "":
model = model.load_from_checkpoint(load_checkpoint_path)
# TODO Train detail params
# Callbacks
val_every_n_epoch = 5
callbacks = [
ModelCheckpoint(period=val_every_n_epoch,
filename=f"{model.generate_model_name()}",
dirpath=save_checkpoint_dir_path,
monitor='val_map',
verbose=True,
result_file_path = os.environ.get("RESULT_FILE_PATH")
checkpoint_path = os.environ.get("CHECKPOINT_PATH")
label_file_path = os.environ.get("LABEL_FILE_PATH")
width, height = int(os.environ.get("IMAGE_WIDTH")), int(
os.environ.get("IMAGE_HEIGHT"))
label_names, label_dict = create_label_list_and_dict(label_file_path)
n_classes = len(label_names)
transforms = AlbumentationsClsWrapperTransform(
A.Compose([
A.Resize(width=width, height=height),
ToTensorV2(),
]))
dataset = ImageOnlyDataset(image_dir=images_dir_path,
image_transform=transforms)
# TODO Choose model, parameters.
print("Loading model...")
model_class = model_service.resolve_classification_model(model_name)
model: ClassificationModel = try_cuda(
model_class.load_from_checkpoint(checkpoint_path))
print("Model loaded")
with open(result_file_path, "w") as file:
file.write(f"filepath,result label\n")
for i, img_tensor in enumerate(tqdm.tqdm(dataset)):
label, prob = model.predict_label(to_4dim(img_tensor))
file.write(f"{dataset.current_file_path()},{label_names[label[0]]}\n")
args = parser.parse_args()
dataset_info = SEGMENTATION_DATASET_INFO.get(args.dataset)
if dataset_info is None:
raise ValueError(f"Invalid dataset name - {args.dataset}. Required [{list(SEGMENTATION_DATASET_INFO.keys())}]")
label_names = dataset_info["label_names"]
class_index_dict = label_names_to_dict(label_names)
# Fetch dataset.
train_transforms, test_transforms = build_transforms(args, dataset_info["n_classes"] + 1)
train_dataset, test_dataset = build_dataset(args, train_transforms, test_transforms)
train_data_loader, test_data_loader = build_data_loader(args, train_dataset, test_dataset)
# Fetch model and load weight.
model = try_cuda(build_model(args, dataset_info["n_classes"] + 1)) # include background class
if args.load_checkpoint_path:
model = model.load_from_checkpoint(args.load_checkpoint_path)
# Training setting.
logger = get_logger("segmentation_demo", args, model)
callbacks = [ModelCheckpoint(period=args.val_every_n_epoch, filename=f"{model.generate_model_name()}",
dirpath=args.save_checkpoint_path, monitor='val_iou', verbose=True, mode="max")]
if args.progress_dir:
callbacks.append(GenerateSegmentationImageCallback(output_dir=args.progress_dir, per_epoch=2, model=model,
dataset=test_dataset))
# Training.
Trainer(max_epochs=args.epoch, callbacks=callbacks, gpus=-1,
check_val_every_n_epoch=args.val_every_n_epoch, logger=logger) \
.fit(model, train_dataloader=train_data_loader, val_dataloaders=test_data_loader)
checkpoint_path = os.environ.get("CHECKPOINT_PATH")
label_file_path = os.environ.get("LABEL_FILE_PATH")
width, height = int(os.environ.get("IMAGE_WIDTH")), int(os.environ.get("IMAGE_HEIGHT"))
if not Path(result_dir_path).exists():
Path(result_dir_path).mkdir()
label_names, label_dict = create_label_list_and_dict(label_file_path)
n_classes = len(label_names)
transforms = AlbumentationsClsWrapperTransform(A.Compose([
A.Resize(width=width, height=height),
ToTensorV2(),
]))
dataset = ImageOnlyDataset(image_dir=images_dir_path, image_transform=transforms)
print("Loading model...")
model_class = model_service.resolve_detection_model(model_name)
model: DetectionModel = try_cuda(model_class.load_from_checkpoint(checkpoint_path))
print("Model loaded")
for i, img_tensor in enumerate(tqdm.tqdm(dataset)):
origin_image = normalize255(tensor_to_cv(img_tensor))
result_bboxes = model.predict_bboxes(to_4dim(img_tensor))[0]
result_img = model.generate_bbox_draw_image(origin_image, bboxes=result_bboxes,
model_img_size=(width, height),
label_names=label_names)
result_img = cv2.resize(result_img, dataset.current_image_size())
cv2.imwrite(f"{result_dir_path}/result_{i}.jpg", result_img)
)
label_names = dataset_info["label_names"]
class_index_dict = label_names_to_dict(label_names)
# Fetch dataset.
train_transforms, test_transforms = build_transforms(
args, class_index_dict)
train_dataset, test_dataset = build_dataset(args, train_transforms,
test_transforms)
train_data_loader, test_data_loader = build_data_loader(
args, train_dataset, test_dataset, AdjustDetectionTensorCollator(),
AdjustDetectionTensorCollator())
# Fetch model and load weight.
model = try_cuda(build_model(args, dataset_info["n_classes"]))
if args.load_checkpoint_path:
model = model.load_from_checkpoint(args.load_checkpoint_path)
# Training setting.
logger = get_logger("detection_demo", args, model)
callbacks = [
ModelCheckpoint(period=args.val_every_n_epoch,
filename=f"{model.generate_model_name()}",
dirpath=args.save_checkpoint_path,
monitor='val_map',
verbose=True,
mode="max")
]
if args.progress_dir:
callbacks.append(
RealtimeDetection
load_dotenv("envs/camera_prediction.env")
model_name = os.environ.get("MODEL_NAME")
checkpoint_path = os.environ.get("CHECKPOINT_PATH")
label_file_path = os.environ.get("LABEL_FILE_PATH")
width, height = int(os.environ.get("IMAGE_WIDTH")), int(
os.environ.get("IMAGE_HEIGHT"))
label_names, label_dict = create_label_list_and_dict(label_file_path)
n_classes = len(label_names)
print("Loading model...")
model_class = model_service.resolve_model_class_from_name(model_name)
model: BaseDeepextModel = try_cuda(
model_class.load_from_checkpoint(checkpoint_path))
print("Model loaded")
if isinstance(model, SegmentationModel):
RealtimeSegmentation(model=model,
img_size_for_model=(width, height)).realtime_predict(
video_output_path="output.mp4")
elif isinstance(model, DetectionModel):
RealtimeDetection(model=model,
img_size_for_model=(width, height),
label_names=label_names).realtime_predict(
video_output_path="output.mp4")
elif isinstance(model, AttentionClassificationModel):
RealtimeAttentionClassification(model=model,
img_size_for_model=(width, height),
label_names=label_names).realtime_predict(
width, height = int(os.environ.get("IMAGE_WIDTH")), int(
os.environ.get("IMAGE_HEIGHT"))
n_classes = int(os.environ.get("N_CLASSES"))
if not Path(result_dir_path).exists():
Path(result_dir_path).mkdir()
transforms = AlbumentationsClsWrapperTransform(
A.Compose([
A.Resize(width=width, height=height),
ToTensorV2(),
]))
dataset = ImageOnlyDataset(image_dir=images_dir_path,
image_transform=transforms)
# TODO Choose model, parameters.
print("Loading model...")
model_class = model_service.resolve_segmentation_model(model_name)
model: SegmentationModel = try_cuda(
model_class.load_from_checkpoint(checkpoint_path))
print("Model loaded")
for i, img_tensor in enumerate(tqdm.tqdm(dataset)):
origin_image = normalize255(tensor_to_cv(img_tensor))
pred_label, prob = model.predict_index_image(to_4dim(img_tensor))
index_image = tensor_to_cv(pred_label[0])
result_img = model.generate_mixed_segment_image(origin_image, index_image)
result_img = cv2.resize(result_img, dataset.current_image_size())
cv2.imwrite(f"{result_dir_path}/result_{i}.jpg", result_img)
