def main():
args = get_args()
torch.distributed.init_process_group(backend="nccl")
with open(args.config_path) as f:
hparams = yaml.load(f, Loader=yaml.SafeLoader)
hparams.update({
"local_rank": args.local_rank,
"fp16": args.fp16,
})
args.output_path.mkdir(parents=True, exist_ok=True)
hparams["output_path"] = args.output_path
device = torch.device("cuda", args.local_rank) # pylint: disable=E1101
model = object_from_dict(hparams["model"])
corrections: Dict[str, str] = {"model.": ""}
state_dict = state_dict_from_disk(file_path=args.weight_path,
rename_in_layers=corrections)
model.load_state_dict(state_dict)
model = nn.Sequential(model, nn.Softmax(dim=1))
model = model.to(device)
if args.fp16:
model = model.half()
model = torch.nn.parallel.DistributedDataParallel(
model, device_ids=[args.local_rank], output_device=args.local_rank)
file_paths = []
for regexp in ["*.jpg", "*.png", "*.jpeg", "*.JPG"]:
file_paths += sorted(args.input_path.rglob(regexp))
# Filter file paths for which we already have predictions
file_paths = [
x for x in file_paths
if not (args.output_path / x.parent.name / f"{x.stem}.txt").exists()
]
dataset = InferenceDataset(file_paths,
transform=from_dict(hparams["test_aug"]))
sampler = DistributedSampler(dataset, shuffle=False)
dataloader = torch.utils.data.DataLoader(
dataset,
batch_size=args.batch_size,
num_workers=args.num_workers,
pin_memory=True,
shuffle=False,
drop_last=False,
sampler=sampler,
)
predict(dataloader, model, hparams, device)
def val_dataloader(self):
val_aug = from_dict(self.hparams["val_aug"])
return DataLoader(
FaceMaskDataset(self.val_samples, val_aug),
batch_size=self.hparams["val_parameters"]["batch_size"],
num_workers=self.hparams["num_workers"],
shuffle=False,
pin_memory=True,
drop_last=False,
)
def val_dataloader(self):
val_aug = from_dict(self.hparams["val_aug"])
result = DataLoader(
SegmentationDataset(self.train_samples, val_aug),
batch_size=self.hparams["val_parameters"]["batch_size"],
num_workers=self.hparams["num_workers"],
shuffle=False,
pin_memory=True,
drop_last=False,
)
return result
def train_dataloader(self):
train_aug = from_dict(self.hparams["train_aug"])
result = DataLoader(
FaceMaskDataset(self.train_samples, train_aug),
batch_size=self.hparams["train_parameters"]["batch_size"],
num_workers=self.hparams["num_workers"],
shuffle=True,
pin_memory=True,
drop_last=True,
)
return result
def val_dataloader(self):
val_aug = from_dict(self.hparams.val_aug)
result = DataLoader(
ClassificationDataset(self.val_image_paths, val_aug, length=None),
batch_size=self.hparams.val_parameters.batch_size,
num_workers=self.hparams.num_workers,
shuffle=False,
pin_memory=True,
drop_last=False,
)
print("Val dataloader = ", len(result))
return result
def val_dataloader(self):
val_aug = from_dict(self.hparams["val_aug"])
result = DataLoader(
SegmentationDataset(self.val_samples, val_aug, length=None),
batch_size=self.hparams["val_parameters"]["batch_size"],
num_workers=self.hparams["num_workers"],
shuffle=False,
pin_memory=True,
drop_last=False,
)
print("Val dataloader = ", len(result))
return result
def val_dataloader(self):
return DataLoader(
FaceDetectionDataset(
label_path=self.hparams["val_annotation_path"],
image_path=self.hparams["val_image_path"],
transform=from_dict(self.hparams["val_aug"]),
preproc=self.preproc,
),
batch_size=self.hparams["val_parameters"]["batch_size"],
num_workers=self.hparams["num_workers"],
shuffle=False,
pin_memory=True,
drop_last=True,
collate_fn=detection_collate,
)
def test_dataloader(self) -> DataLoader:
return DataLoader(
InferenceDataset(
self.hparams["file_paths"],
origin_size=self.hparams["origin_size"],
target_size=self.hparams["target_size"],
max_size=self.hparams["max_size"],
transform=from_dict(self.hparams["test_aug"]),
),
batch_size=1,
num_workers=self.hparams["num_workers"],
shuffle=False,
pin_memory=True,
drop_last=False,
)
def val_dataloader(self):
val_aug = from_dict(self.hparams["val_aug"])
result = DataLoader(
SegmentationDataset(self.val_samples, val_aug, length=None),
batch_size=self.hparams["val_parameters"]["batch_size"],
num_workers=self.hparams["num_workers"],
shuffle=False,
pin_memory=True,
drop_last=False,
)
print("Val dataloader = ", len(result))
# self.logger.experiment.log({"val_input_image": [wandb.Image(result["mask"].cpu(), caption="val_input_image")]})
return result
def val_dataloader(self) -> DataLoader:
result = DataLoader(
FaceDetectionDataset(
label_path=VAL_LABEL_PATH,
image_path=VAL_IMAGE_PATH,
transform=from_dict(self.config.val_aug),
preproc=self.preproc,
rotate90=self.config.val_parameters.rotate90,
),
batch_size=self.config.val_parameters.batch_size,
num_workers=self.config.num_workers,
shuffle=False,
pin_memory=True,
drop_last=True,
collate_fn=detection_collate,
)
return result
def train_dataloader(self):
result = DataLoader(
FaceDetectionDataset(
label_path=TRAIN_LABEL_PATH,
image_path=TRAIN_IMAGE_PATH,
transform=from_dict(self.config.train_aug),
preproc=self.preproc,
rotate90=self.config.train_parameters.rotate90,
),
batch_size=self.config.train_parameters.batch_size,
num_workers=self.config.num_workers,
shuffle=True,
pin_memory=True,
drop_last=False,
collate_fn=detection_collate,
)
print("Len train dataloader = ", len(result))
return result
def train_dataloader(self):
train_aug = from_dict(self.hparams["train_aug"])
if "epoch_length" not in self.hparams["train_parameters"]:
epoch_length = None
else:
epoch_length = self.hparams["train_parameters"]["epoch_length"]
result = DataLoader(
SegmentationDataset(self.train_samples, train_aug, epoch_length),
batch_size=self.hparams["train_parameters"]["batch_size"],
num_workers=self.hparams["num_workers"],
shuffle=True,
pin_memory=True,
drop_last=True,
)
print("Train dataloader = ", len(result))
return result
def train_dataloader(self):
train_aug = from_dict(self.hparams.train_aug)
if "epoch_length" not in self.hparams.train_parameters:
epoch_length = None
else:
epoch_length = self.hparams.train_parameters.epoch_length
result = DataLoader(
ClassificationDataset(self.train_image_paths, train_aug,
epoch_length),
batch_size=self.hparams.train_parameters.batch_size,
num_workers=self.hparams.num_workers,
shuffle=True,
pin_memory=True,
drop_last=True,
)
print("Train dataloader = ", len(result))
return result
def dataloaders(self, **kwargs):
dataloaders = {}
for type, data in self.datasets.items():
if "dataset" not in data or "dataloader" not in data:
raise ValueError("dataset info missing for {%type}")
dataset_args = {}
if "transforms" in data:
dataset_args.update(
{"transforms": from_dict(data["transform"])})
else:
dataset_args.update({
"transforms":
kwargs["transforms"][type] if "transforms" in kwargs
and type in kwargs["transforms"] else None
})
if "folds" in kwargs and type in kwargs["folds"]:
dataset_args.update({"fold": kwargs["folds"][type]})
elif "folds" in data:
dataset_args.update({"fold": data["folds"]})
dataset = create_dataset(data["dataset"]["class"], **dataset_args,
**data["dataset"].get("kwargs", {}))
if "sampler" in data:
sampler = create_sampler(data["sampler"]["class"], dataset,
data["sampler"].get("kwargs", {}))
else:
sampler = None
if "collate_fn" in data:
collate = get_collator(data["collate_fn"])
else:
collate = None
dataloader = create_dataloader(
data["dataloader"]["class"], dataset, sampler, collate,
**data["dataloader"].get("kwargs", {}))
dataloaders.update({type: dataloader})
return dataloaders
def main():
args = get_args()
# torch.distributed.init_process_group(backend="nccl", rank=args.local_rank)
with open(args.config_path) as f:
hparams = yaml.load(f, Loader=yaml.SafeLoader)
hparams['model']['pretrained'] = False
hparams.update({
"json_path":
args.output_path,
"visualize":
args.visualize,
"confidence_threshold":
args.confidence_threshold,
"nms_threshold":
args.nms_threshold,
"keep_top_k":
args.keep_top_k,
"local_rank":
args.local_rank,
"prior_box":
object_from_dict(hparams["prior_box"],
image_size=[args.max_size, args.max_size]),
"fp16":
args.fp16,
})
if args.visualize:
output_vis_path = args.output_path / "viz"
output_vis_path.mkdir(parents=True, exist_ok=True)
hparams["output_vis_path"] = output_vis_path
output_label_path = args.output_path / "labels"
output_label_path.mkdir(parents=True, exist_ok=True)
hparams["output_label_path"] = output_label_path
device = torch.device('cuda')
# device = torch.device("cuda", args.local_rank)
model = object_from_dict(hparams["model"])
model = model.to(device)
if args.fp16:
model = model.half()
corrections: Dict[str, str] = {"model.": ""}
checkpoint = load_checkpoint(file_path=args.weight_path,
rename_in_layers=corrections)
model.load_state_dict(checkpoint["state_dict"])
# model = torch.nn.parallel.DistributedDataParallel(
# model, device_ids=[args.local_rank], output_device=args.local_rank
# )
file_paths = sorted([x for x in args.input_path.rglob("*") if x.is_file()])
dataset = InferenceDataset(file_paths,
max_size=args.max_size,
transform=from_dict(hparams["test_aug"]))
# sampler = DistributedSampler(dataset, shuffle=False)
dataloader = torch.utils.data.DataLoader(
dataset,
batch_size=args.batch_size,
num_workers=args.num_workers,
pin_memory=True,
shuffle=False,
drop_last=False,
# sampler=sampler,
)
predict(dataloader, model, hparams, device)
def main():
args = get_args()
torch.distributed.init_process_group(backend="nccl")
with open(args.config_path) as f:
hparams = yaml.load(f, Loader=yaml.SafeLoader)
hparams.update({
"json_path": args.output_path,
"local_rank": args.local_rank,
"fp16": args.fp16
})
device = torch.device("cuda", args.local_rank)
model = object_from_dict(hparams["model"])
model = model.to(device)
if args.fp16:
model = model.half()
corrections: Dict[str, str] = {"model.": ""}
checkpoint = load_checkpoint(file_path=args.weight_path,
rename_in_layers=corrections)
model.load_state_dict(checkpoint["state_dict"])
model = torch.nn.parallel.DistributedDataParallel(
model, device_ids=[args.local_rank], output_device=args.local_rank)
file_paths = []
for regexp in ["*.jpg", "*.png", "*.jpeg", "*.JPG"]:
file_paths += sorted(x for x in tqdm(args.input_path.rglob(regexp)))
dataset = FaceMaskTestDataset(file_paths,
transform=from_dict(hparams["test_aug"]))
sampler = DistributedSampler(dataset, shuffle=False)
dataloader = torch.utils.data.DataLoader(
dataset,
batch_size=args.batch_size,
num_workers=args.num_workers,
pin_memory=True,
shuffle=sampler is None,
drop_last=False,
sampler=sampler,
)
prediction = predict(dataloader, model, hparams, device)
prediction_list = [
torch.zeros_like(prediction) for _ in range(dist.get_world_size())
]
dist.all_gather(prediction_list, prediction)
if dist.get_rank() == 0:
with torch.no_grad():
predictions = torch.cat(
prediction_list,
dim=1).reshape(-1).cpu().numpy()[:len(dataset)]
df = pd.DataFrame({
"file_path": file_paths,
"predictions": predictions
})
df.to_csv(args.output_path, index=False)
def main() -> None:
args = get_args()
torch.distributed.init_process_group(backend="nccl")
with args.config_path.open() as f:
hparams = yaml.load(f, Loader=yaml.SafeLoader)
hparams.update(
{
"json_path": args.output_path,
"visualize": args.visualize,
"confidence_threshold": args.confidence_threshold,
"nms_threshold": args.nms_threshold,
"keep_top_k": args.keep_top_k,
"local_rank": args.local_rank,
"prior_box": object_from_dict(hparams["prior_box"], image_size=[args.max_size, args.max_size]),
"fp16": args.fp16,
"folder_in_name": args.folder_in_name,
}
)
if args.visualize:
output_vis_path = args.output_path / "viz"
output_vis_path.mkdir(parents=True, exist_ok=True)
hparams["output_vis_path"] = output_vis_path
output_label_path = args.output_path / "labels"
output_label_path.mkdir(parents=True, exist_ok=True)
hparams["output_label_path"] = output_label_path
device = torch.device("cuda", args.local_rank)
model = object_from_dict(hparams["model"])
model = model.to(device)
if args.fp16:
model = model.half()
corrections: Dict[str, str] = {"model.": ""}
state_dict = state_dict_from_disk(file_path=args.weight_path, rename_in_layers=corrections)
model.load_state_dict(state_dict)
model = torch.nn.parallel.DistributedDataParallel(
model, device_ids=[args.local_rank], output_device=args.local_rank
)
file_paths = list(args.input_path.rglob("*.jpg"))
dataset = InferenceDataset(file_paths, max_size=args.max_size, transform=from_dict(hparams["test_aug"]))
sampler: DistributedSampler = DistributedSampler(dataset, shuffle=False)
dataloader = torch.utils.data.DataLoader(
dataset,
batch_size=args.batch_size,
num_workers=args.num_workers,
pin_memory=True,
shuffle=False,
drop_last=False,
sampler=sampler,
)
predict(dataloader, model, hparams, device)
def main():
args = get_args()
with open(args.config_path) as f:
hparams = yaml.load(f, Loader=yaml.SafeLoader)
model = SegmentDocs(hparams)
test_file_names = sorted(Path(args.image_path).glob("*.jpg"))
test_mask_path = args.output_path / "masks"
test_vis_path = args.output_path / "vis"
test_mask_path.mkdir(exist_ok=True, parents=True)
test_vis_path.mkdir(exist_ok=True, parents=True)
test_aug = from_dict(hparams["test_aug"])
dataloader = DataLoader(
SegmentationDatasetTest(test_file_names, test_aug),
batch_size=hparams["test_parameters"]["batch_size"],
num_workers=hparams["num_workers"],
shuffle=False,
pin_memory=True,
drop_last=False,
)
corrections: Dict[str, str] = {}
checkpoint = load_checkpoint(file_path=args.checkpoint_path, rename_in_layers=corrections) # type: ignore
model.load_state_dict(checkpoint["state_dict"])
model = nn.Sequential(model, nn.Sigmoid())
model = tta.MultiscaleTTAWrapper(model, [0.5, 2])
model.eval()
model = model.half()
model.cuda()
with torch.no_grad():
for batch in tqdm(dataloader):
features = batch["features"]
image_ids = batch["image_id"]
preds = tta.fliplr_image2mask(model, features.half().cuda())
for batch_id in range(features.shape[0]):
image_id = image_ids[batch_id]
mask = (preds[batch_id][0] > 0.5).cpu().numpy().astype(np.uint8)
height = batch["height"][batch_id].item()
width = batch["width"][batch_id].item()
pads = batch["pads"][batch_id].cpu().numpy()
mask = unpad(mask, pads)
mask = remove_small_connected_binary(mask, min_area=100)
mask = fill_small_holes(mask, min_area=100)
mask = cv2.resize(
mask, (width, height), interpolation=cv2.INTER_NEAREST
)
cv2.imwrite(str(test_mask_path / f"{image_id}.png"), mask * 255)
image = cv2.imread(str(args.image_path / f"{image_id}.jpg"))
mask_image = mask_overlay(image, mask)
cv2.imwrite(
str(test_vis_path / f"{image_id}.jpg"),
np.hstack(
[mask_image, cv2.cvtColor((mask * 255), cv2.COLOR_GRAY2BGR)]
),
)
"interval": self.hparams["scheduler"]["interval"],
"frequency": self.hparams["scheduler"]["frequency"],
}
return [optimizer], [scheduler_dict]
if __name__ == "__main__":
args, _ = _parse_args()
with open(args.config_path) as f:
hparams = yaml.load(f, Loader=yaml.SafeLoader)
hparams["model"]["model"]["classes"] = len(hparams["categories"])
pl.seed_everything(hparams["seed"])
hparams["data"]["data"]["transforms"] = {
"train": from_dict(hparams["data"]["transforms"]["train"]),
"val": from_dict(hparams["data"]["transforms"]["val"]),
"test": from_dict(hparams["data"]["transforms"]["test"]),
}
data = SegmentationDataModule(**hparams["data"]["data"])
model = SegmentationModule(hparams["model"])
trainer = object_from_dict(
hparams["trainer"]["trainer"],
logger=object_from_dict(hparams["trainer"]["logger"]),
callbacks=[
object_from_dict(callback)
for callback in hparams["trainer"]["callbacks"].values()
],
)
trainer.fit(model, data)
trainer.test()
def load_transforms(transforms: dict):
return from_dict(transforms)
def main():
args = get_args()
torch.set_grad_enabled(False)
with open(args.config_path) as f:
hparams = yaml.load(f, Loader=yaml.SafeLoader)
device = torch.device("cpu" if args.cpu else "cuda")
net = get_model(hparams, args.weights, args.fp16, device)
file_paths = sorted(args.input_path.rglob("*.jpg"))
output_path = args.output_path
output_vis_path = output_path / "viz"
output_label_path = output_path / "labels"
output_image_path = output_path / "images"
prepare_output_folders(output_vis_path, output_label_path,
output_image_path, args.save_boxes, args.save_crops,
args.visualize)
transform = from_dict(hparams["test_aug"])
test_loader = DataLoader(
InferenceDataset(file_paths, args.origin_size, transform=transform),
batch_size=args.batch_size,
num_workers=args.num_workers,
pin_memory=True,
drop_last=False,
)
with torch.no_grad():
for raw_input in tqdm(test_loader):
torched_images = raw_input["torched_image"].type(net.dtype)
resizes = raw_input["resize"]
image_paths = raw_input["image_path"]
raw_images = raw_input["raw_image"]
labels = []
if (args.batch_size == 1 and args.save_boxes
and (output_label_path /
f"{Path(image_paths[0]).stem}.json").exists()):
continue
loc, conf, land = net(torched_images.to(device)) # forward pass
batch_size = torched_images.shape[0]
image_height, image_width = torched_images.shape[2:]
scale1 = torch.Tensor([
image_width,
image_height,
image_width,
image_height,
image_width,
image_height,
image_width,
image_height,
image_width,
image_height,
])
scale1 = scale1.to(device)
scale = torch.Tensor(
[image_width, image_height, image_width, image_height])
scale = scale.to(device)
priors = object_from_dict(hparams["prior_box"],
image_size=(image_height,
image_width)).to(loc.device)
for batch_id in range(batch_size):
image_path = image_paths[batch_id]
file_id = Path(image_path).stem
raw_image = raw_images[batch_id]
resize = resizes[batch_id].float()
boxes = decode(loc.data[batch_id], priors,
hparams["test_parameters"]["variance"])
boxes *= scale / resize
scores = conf[batch_id][:, 1]
landmarks = decode_landm(
land.data[batch_id], priors,
hparams["test_parameters"]["variance"])
landmarks *= scale1 / resize
# ignore low scores
valid_index = torch.where(
scores > args.confidence_threshold)[0]
boxes = boxes[valid_index]
landmarks = landmarks[valid_index]
scores = scores[valid_index]
order = scores.argsort(descending=True)
boxes = boxes[order]
landmarks = landmarks[order]
scores = scores[order]
# do NMS
keep = nms(boxes, scores, args.nms_threshold)
boxes = boxes[keep, :].int()
landmarks = landmarks[keep].int()
if boxes.shape[0] == 0:
continue
scores = scores[keep].cpu().numpy().astype(np.float64)
if args.visualize:
vis_image = raw_image.cpu().numpy().copy()
for crop_id, bbox in enumerate(boxes):
bbox = bbox.cpu().numpy()
labels += [{
"crop_id": crop_id,
"bbox": bbox.tolist(),
"score": scores[crop_id],
"landmarks": landmarks[crop_id].tolist(),
}]
if args.save_crops:
x_min, y_min, x_max, y_max = bbox
x_min = np.clip(x_min, 0, x_max - 1)
y_min = np.clip(y_min, 0, y_max - 1)
crop = raw_image[y_min:y_max,
x_min:x_max].cpu().numpy()
if args.visualize:
vis_image = cv2.rectangle(vis_image,
(x_min, y_min),
(x_max, y_max),
color=(255, 0, 0),
thickness=3)
target_folder = output_image_path / f"{file_id}"
target_folder.mkdir(exist_ok=True, parents=True)
crop_file_path = target_folder / f"{file_id}_{crop_id}.jpg"
if crop_file_path.exists():
continue
cv2.imwrite(str(crop_file_path),
cv2.cvtColor(crop, cv2.COLOR_BGR2RGB))
if args.visualize:
cv2.imwrite(str(output_vis_path / f"{file_id}.jpg"),
cv2.cvtColor(vis_image, cv2.COLOR_BGR2RGB))
if args.save_boxes:
result = {
"file_path": image_path,
"file_id": file_id,
"bboxes": labels,
}
with open(output_label_path / f"{file_id}.json", "w") as f:
json.dump(result, f, indent=2)
