def predict_single():
device = "cpu"
image = np.array(Image.open('data/val_images/0.png').convert("RGB"))
mask = np.array(Image.open('data/val_masks/0.bmp').convert("L"), dtype=np.float32)
mask[mask == 255.0] = 1
augmentations = val_transforms(image=image, mask=mask)
image = augmentations["image"]
mask = augmentations["mask"]
plt.imshow(image.squeeze().permute(1, 2, 0))
plt.show()
plt.imshow(mask, cmap='gray')
plt.show()
image = torch.tensor(image, requires_grad=True).to(device)
image = image.unsqueeze(0)
model = UNET(in_channels=3, out_channels=1).to(device)
load_checkpoint(torch.load("check_Unet_99_95.pth.tar", map_location=torch.device('cpu')), model)
# image = image.to(device=device)
model.eval()
with torch.no_grad():
preds = torch.sigmoid(model(image))
preds = (preds > 0.5).float()
torchvision.utils.save_image(preds, "./pred_100.png")
model.train()
def main():
model = UNET(in_channels=3, out_channels=1).to(device)
loss_fun = nn.BCEWithLogitsLoss()
optimizer = optim.Adam(model.parameters(), lr=lr)
train_loader, test_loader = get_loaders(batch_size)
print("Training Model")
print("==============")
epoch_count = 1
for epoch in range(epochs):
print("Epoch ", epoch_count)
print("---------")
train_loading_bar = tqdm(train_loader, position=0, leave=True)
model.train()
train_correct_pixels = 0
train_total_pixels = 0
count = 0
# iterate over the train data loader
for _, (pixel_data, target_masks) in enumerate(train_loading_bar):
count += 1
pixel_data = pixel_data.to(device=device)
target_masks_unsqueezed = target_masks.float().unsqueeze(1).to(
device=device)
model.zero_grad()
predictions = model(pixel_data)
loss = loss_fun(predictions, target_masks_unsqueezed)
loss.backward()
# get and accumualate the train accuracy
(correct_pixels,
total_pixels) = get_accuracy(predictions, target_masks, device)
train_correct_pixels = train_correct_pixels + correct_pixels
train_total_pixels = train_total_pixels + total_pixels
optimizer.step()
train_loading_bar.set_postfix(loss=loss.item())
print(
f"\nTrain Accuracy: {train_correct_pixels/train_total_pixels*100:.2f}%"
)
model.eval()
epoch_count += 1
# save model upon training
print("Training Complete!")
Path(my_path + "/model").mkdir(parents=True, exist_ok=True)
torch.save(model.state_dict(), r"model" + r"\blueno_detection.pth")
test_loading_bar = tqdm(test_loader)
test_correct_pixels = 0
test_total_pixels = 0
print("Testing Model")
print("=============")
count = 0
# iterate over the test data loader
for _, (pixel_data, target_masks) in enumerate(test_loading_bar):
count += 1
pixel_data = pixel_data.to(device=device)
target_masks_unsqueezed = target_masks.float().unsqueeze(1).to(
device=device)
predictions = model(pixel_data)
# get and accumualate the test accuracy
(correct_pixels, total_pixels) = get_accuracy(predictions,
target_masks, device)
test_correct_pixels = test_correct_pixels + correct_pixels
test_total_pixels = test_total_pixels + total_pixels
test_loading_bar.set_postfix(loss=loss.item())
print(f"\nTest Accuracy: {test_correct_pixels/test_total_pixels*100:.2f}%")
device = "cpu"
image = np.array(Image.open('19.png').convert("RGB"))
mask = np.array(Image.open('19.bmp').convert("L"), dtype=np.float32)
mask[mask == 255.0] = 1
augmentations = val_transforms(image=image, mask=mask)
image = augmentations["image"]
mask = augmentations["mask"]
plt.imshow(image.squeeze().permute(1, 2, 0))
plt.show()
plt.imshow(mask, cmap='gray')
plt.show()
image = torch.tensor(image, requires_grad=True).to(device)
image = image.unsqueeze(0)
model = UNET(in_channels=3, out_channels=1).to(device)
print("=> Loading checkpoint")
model.load_state_dict(
torch.load("check_Unet_99_95.pth.tar",
map_location=torch.device('cpu'))["state_dict"])
# image = image.to(device=device)
model.eval()
with torch.no_grad():
preds = torch.sigmoid(model(image))
preds = (preds > 0.5).float()
torchvision.utils.save_image(preds, "./pred_100.png")
model.train()