import torch
import torch.nn as nn
import torch.nn.functional as Ftorch
from torch.utils.data import DataLoader
from tqdm import *
from models import *
from augmentation import *
from dataset import *
import glob
device = torch.device('cuda')
def predict(model, loader):
model.eval()
preds = []
with torch.no_grad():
for dct in tqdm(loader, total=len(loader)):
images = dct['images'].to(device)
# meta = dct["meta"].to(device)
pred = model(images)
pred = Ftorch.sigmoid(pred)
pred = pred.detach().cpu().numpy()
preds.append(pred)
preds = np.concatenate(preds, axis=0)
return preds
def get_best_checkpoints(checkpoint_dir, n_best=3, minimize_metric=True):
files = glob.glob(f"{checkpoint_dir}/checkpoints/best*.pth")
files = [file for file in files if not 'full' in file]
top_best_metrics = []
for file in files:
ckp = torch.load(file)
valid_metric = ckp['valid_metrics']['loss']
top_best_metrics.append((file, valid_metric))
top_best_metrics = sorted(
top_best_metrics,
key=lambda x: x[1],
reverse=not minimize_metric
)
top_best_metrics = top_best_metrics[:n_best]
return top_best_metrics
def predict_test_tta_ckp():
test_csv = "./csv/patient2_kfold/test.csv"
# test_root = "/data/stage_1_test_3w/"
# test_root = "/data/png/test_stage_1/adjacent-brain-cropped/"
# test_root = "/data/stage_1_test_3w/"
test_root = "/data/stage_1_test_images_jpg_preprocessing/"
image_type = 'jpg'
image_size = [512, 512]
backbone = "densenet169"
normalization = True
# fold = 2
for fold in [0, 1, 2, 3, 4]:
# /logs/rsna/test/resnet50-anju-512-resume-0/checkpoints//train512.13.pth
scheme = f"{backbone}-mw-512-resume-{fold}"
log_dir = f"/logs/rsna/test/{scheme}/"
with_any = True
if with_any:
num_classes = 6
target_cols = LABEL_COLS
else:
num_classes = 5
target_cols = LABEL_COLS_WITHOUT_ANY
# test_preds = 0
top_best_metrics = get_best_checkpoints(log_dir, n_best=1, minimize_metric=True)
test_preds = 0
for best_metric in top_best_metrics:
checkpoint_path, checkpoint_metric = best_metric
print("*" * 50)
print(f"checkpoint: {checkpoint_path}")
print(f"Metric: {checkpoint_metric}")
model = CNNFinetuneModels(
model_name=backbone,
num_classes=num_classes,
pretrained=False
)
ckp = os.path.join(log_dir, f"checkpoints/best.pth")
checkpoint = torch.load(ckp)
model.load_state_dict(checkpoint['model_state_dict'])
model = nn.DataParallel(model)
model = model.to(device)
augs = test_tta(image_size, normalization)
for name, aug in augs.items():
print("Augmentation: {}".format(name))
test_dataset = RSNADataset(
csv_file=test_csv,
root=test_root,
with_any=with_any,
transform=aug,
mode="test",
image_type=image_type
)
test_loader = DataLoader(
dataset=test_dataset,
batch_size=64,
shuffle=False,
num_workers=8,
)
test_preds += predict(model, test_loader) / (len(augs) * len(top_best_metrics))
os.makedirs(f"/logs/prediction/{scheme}", exist_ok=True)
np.save(f"/logs/prediction/{scheme}/test_{fold}_ckp_tta.npy", test_preds)
test_df = pd.read_csv(test_csv)
test_ids = test_df['sop_instance_uid'].values
ids = []
labels = []
for i, id in enumerate(test_ids):
if not "ID" in id:
id = "ID_" + id
pred = test_preds[i]
for j, target in enumerate(target_cols):
id_target = id + "_" + target
ids.append(id_target)
labels.append(pred[j])
if not with_any:
id_target = id + "_" + "any"
ids.append(id_target)
labels.append(pred.max())
submission_df = pd.DataFrame({
'ID': ids,
'Label': labels
})
submission_df.to_csv(f"/logs/prediction/{scheme}/{scheme}_ckp_tta.csv", index=False)
def predict_valid_tta_ckp():
# test_root = "/data/png/train/adjacent-brain-cropped/"
# test_root = "/data/stage_1_train_3w/"
test_root = "/data/stage_1_test_images_jpg_preprocessing/"
image_type = 'jpg'
image_size = [512, 512]
backbone = "densenet169"
normalization = True
# fold = 2
for fold in [0, 1, 2, 3, 4]:
test_csv = f"./csv/patient2_kfold/valid_{fold}.csv"
# /logs/rsna/test/resnet50-anju-512-resume-0/checkpoints//train512.13.pth
scheme = f"{backbone}-mw-512-resume-{fold}"
log_dir = f"/logs/rsna/test/{scheme}/"
with_any = True
if with_any:
num_classes = 6
target_cols = LABEL_COLS
else:
num_classes = 5
target_cols = LABEL_COLS_WITHOUT_ANY
# test_preds = 0
top_best_metrics = get_best_checkpoints(log_dir, n_best=1, minimize_metric=True)
test_preds = 0
for best_metric in top_best_metrics:
checkpoint_path, checkpoint_metric = best_metric
print("*" * 50)
print(f"checkpoint: {checkpoint_path}")
print(f"Metric: {checkpoint_metric}")
model = CNNFinetuneModels(
model_name=backbone,
num_classes=num_classes,
)
ckp = os.path.join(log_dir, f"checkpoints/best.pth")
checkpoint = torch.load(ckp)
model.load_state_dict(checkpoint['model_state_dict'])
model = nn.DataParallel(model)
model = model.to(device)
augs = test_tta(image_size, normalization)
for name, aug in augs.items():
print("Augmentation: {}".format(name))
test_dataset = RSNADataset(
csv_file=test_csv,
root=test_root,
with_any=with_any,
transform=aug,
mode="valid",
image_type=image_type
)
test_loader = DataLoader(
dataset=test_dataset,
batch_size=64,
shuffle=False,
num_workers=8,
)
test_preds += predict(model, test_loader) / (len(augs) * len(top_best_metrics))
os.makedirs(f"/logs/prediction/{scheme}", exist_ok=True)
np.save(f"/logs/prediction/{scheme}/valid_{scheme}.npy", test_preds)
test_df = pd.read_csv(test_csv)
test_ids = test_df['sop_instance_uid'].values
ids = []
labels = []
for i, id in enumerate(test_ids):
if not "ID" in id:
id = "ID_" + id
pred = test_preds[i]
for j, target in enumerate(target_cols):
id_target = id + "_" + target
ids.append(id_target)
labels.append(pred[j])
if not with_any:
id_target = id + "_" + "any"
ids.append(id_target)
labels.append(pred.max())
submission_df = pd.DataFrame({
'ID': ids,
'Label': labels
})
submission_df.to_csv(f"/logs/prediction/{scheme}/valid_{scheme}.csv", index=False)
if __name__ == '__main__':
# predict_test()
predict_test_tta_ckp()
# predict_valid_tta_ckp()
Datasets
Models
