def setup_model_dataloader(args, batch_size):
num_class = 1
model = ResNetUNet(num_class)
device = torch.device("cuda")
model_file = os.path.join(args.model_dir, 'best_model.pth')
model.load_state_dict(torch.load(model_file, map_location="cuda:0"))
model.to(device)
# Set model to the evaluation mode
model.eval()
# Setup dataset
# Need to be careful here. This isn't perfect.
# I'm assuming that the dataset isn't changing between training and inference time
bee_ds = BeePointDataset(root_dir=args.data_dir)
if 1:
dbfile = open(os.path.join(args.model_dir, 'test_ds.pkl'), 'rb')
test_ds = pickle.load(dbfile)
#test_loader = DataLoader(test_ds, batch_size=1, shuffle=False, num_workers=1)
test_loader = DataLoader(test_ds,
batch_size=batch_size,
shuffle=True,
num_workers=1,
collate_fn=helper.bee_collate_fn)
else:
# Just use the defaults in the bee_ds
test_loader = DataLoader(bee_ds,
batch_size=batch_size,
shuffle=True,
num_workers=1,
collate_fn=helper.bee_collate_fn)
return model, test_loader, device
def setup_model_dataloader(args, batch_size):
num_class = 1
model = ResNetUNet(num_class)
device = torch.device("cuda")
model_file = os.path.join(args.model_dir, 'best_model.pth')
model.load_state_dict(torch.load(model_file, map_location="cuda:0"))
model.to(device)
# Set model to the evaluation mode
model.eval()
return model, device
def predict(weights_path, img):
"""
Args:
weights_path (str): path to where the weights are stored, which will be loaded unto the model
img (PIL): Nuclei PIL image. Image size = 1000x1000
Returns:
Mask (PIL): PIL image detailing the predicted segmentation of the nuclei
"""
# Prepare image
transform = transforms.Compose([
transforms.Pad(
12), # given image is 1000x1000, pad it to make it 1024x1024
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406],
[0.229, 0.224, 0.225]) # imagenet normalization
])
img = transform(img)
img = torch.unsqueeze(img, dim=0)
# Prepare model
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
num_class = 3
model = ResNetUNet(num_class).to(device)
model.load_state_dict(torch.load(weights_path))
model.eval()
# Predict
with torch.no_grad():
outputs = torch.sigmoid(model(img.to(device)))
pred = outputs.to('cpu').detach().numpy()[0].transpose((1, 2, 0))
mask = Image.fromarray(np.uint8(pred * 255)).convert('RGB')
return mask
bboxes = restore_bboxes(cls, rho, theta)
bboxes[:, :8] = bboxes[:, :8] * 4 * (
w / img.width + h / img.height) / 2 # restore scale and resize
return bboxes
if __name__ == '__main__':
img_files = [
os.path.join(cfg.test.dataset_test, img_file)
for img_file in sorted(os.listdir(cfg.test.dataset_test))
]
img_path = np.random.choice(img_files)
print(img_path)
model_path = cfg.test.model_pth
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model = ResNetUNet(backbone='50')
model.load_state_dict(torch.load(model_path))
model.to(device)
model.eval()
img = Image.open(img_path).convert('RGB')
img_array = np.array(img)
bboxes = detect_single_image(img, model, device)
print(bboxes.shape)
for bbox in bboxes:
pts = bbox[:8].astype(np.int32).reshape((-1, 1, 2))
cv2.polylines(img_array, [pts], True, (0, 215, 255), 3)
Image.fromarray(img_array).save('res.png')