test_images_dir = os.environ.get("TEST_IMAGES_PATH")
test_annotations_dir = os.environ.get("TEST_ANNOTATIONS_PATH")
label_file_path = os.environ.get("LABEL_FILE_PATH")
load_checkpoint_path = os.environ.get("LOAD_CHECKPOINT_PATH")
save_checkpoint_dir_path = os.environ.get("SAVE_CHECKPOINT_DIR_PATH")
progress_dir = os.environ.get("PROGRESS_DIR_PATH")
# Model params
width, height = int(os.environ.get("IMAGE_WIDTH")), int(
os.environ.get("IMAGE_HEIGHT"))
# Learning params
batch_size = int(os.environ.get("BATCH_SIZE"))
lr = float(os.environ.get("LR"))
epoch = int(os.environ.get("EPOCH"))
label_names, class_index_dict = create_label_list_and_dict(label_file_path)
n_classes = len(label_names)
ignore_indices = [
255,
]
# TODO Data augmentation
train_transforms = AlbumentationsDetectionWrapperTransform([
A.HorizontalFlip(),
A.RandomResizedCrop(width=width, height=height, scale=(0.8, 1.)),
A.OneOf([
A.Blur(blur_limit=5),
A.RandomBrightnessContrast(),
A.RandomGamma(),
]),
from deepext_with_lightning.image_process.convert import try_cuda, normalize255, tensor_to_cv, to_4dim
from deepext_with_lightning.dataset.functions import create_label_list_and_dict
load_dotenv("envs/attention_classification.env")
images_dir_path = os.environ.get("IMAGES_DIR_PATH")
result_dir_path = os.environ.get("RESULT_DIR_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"))
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: AttentionClassificationModel = try_cuda(
AttentionBranchNetwork(
n_classes=n_classes,
args = parser.parse_args()
# Fetch model and load weight.
model_class = model_service.resolve_model_class_from_name(args.model)
model: BaseDeepextModel = try_cuda(
model_class.load_from_checkpoint(args.load_checkpoint_path))
if isinstance(model, SegmentationModel):
RealtimeSegmentation(
model=model,
img_size_for_model=(args.image_size,
args.image_size)).realtime_predict(
video_output_path="output.mp4")
elif isinstance(model, DetectionModel):
label_names, label_dict = create_label_list_and_dict(
args.label_names_path)
RealtimeDetection(model=model,
img_size_for_model=(args.image_size,
args.image_size),
label_names=label_names).realtime_predict(
video_output_path="output.mp4")
if isinstance(model, AttentionClassificationModel):
label_names, label_dict = create_label_list_and_dict(
args.label_names_path)
RealtimeAttentionClassification(
model=model,
img_size_for_model=(args.image_size, args.image_size),
label_names=label_names).realtime_predict(
video_output_path="output.mp4")
elif isinstance(model, ClassificationModel):
label_names, label_dict = create_label_list_and_dict(