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)