def forward(ctx, points, h_value, device, d, out_size, version):
if version == 2:
heat_map = utils.heat_map_tensor(points.view(-1, 2),
h_value.view(-1, 1), device, d,
out_size)
if version == 1 or version == 0:
heat_map = utils.heat_map_tensor(points, h_value, device, d,
out_size)
ctx.save_for_backward(points, h_value, heat_map)
ctx.out_size = out_size
ctx.device = device
ctx.version = version
return heat_map
def test_model(image_path, data, model, arg):
device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
img_transform = utils.transform(arg.version)
model.train(False)
_, _, _, stain, _ = steinseperation.stainsep(image_path)
H_Channel = stain[0]
cropped, coords = utils.patch_extraction(data,
H_Channel,
arg.patch_size,
arg.stride_size,
arg.version)
H, W, _ = data.shape
cell = np.zeros((H, W))
count = np.zeros((H, W))
[H_prime, W_prime] = arg.heatmap_size
for img, coord in zip(cropped, coords):
img = img_transform(np.uint8(img))
img = img.float()
img = img.to(device)
img = img.unsqueeze(0)
point, h = model(img)
heat_map = utils.heat_map_tensor(point.view(-1, 2),
h.view(-1, 1),
device,
arg.d,
arg.heatmap_size)
heatmap = heat_map.cpu().detach().numpy().reshape((H_prime, W_prime))
start_H, end_H, start_W, end_W = utils.find_out_coords(coord,
arg.patch_size,
arg.heatmap_size)
cell[start_H:end_H, start_W:end_W] += heatmap
idx = np.argwhere(heatmap != 0)
count[idx[:,0]+start_H, idx[:,1]+start_W] += 1
count[count==0] = 1
cell = np.divide(cell, count)
return cell
def print_values(img, model, arg):
device = torch.device(
'cuda') if torch.cuda.is_available() else torch.device('cpu')
img_transform = utils.transform(arg.version)
img = img_transform(np.uint8(img))
img = img.float()
img = img.to(device)
img = img.unsqueeze(0)
model.train(False)
point, h = model(img)
print('points :', point.cpu().detach().numpy().reshape((-1, 2)))
print('H :', h.cpu().detach().numpy().reshape((-1, 1)))
print(20 * '*')
return utils.heat_map_tensor(point.view(-1, 2), h.view(-1, 1), device,
arg.d, arg.heatmap_size)
def detect_cell(self, file):
dataset = self.data_loader(file)
H, W = dataset.dimension()
cell = np.zeros((H, W))
count = np.zeros((H, W))
[H_prime, W_prime] = self.heatmap_size
for data in dataset:
img, coords = data
img = img.to(self.device)
img = img.unsqueeze(0)
point, h = self.SC_CNN.model(img)
heat_map = utils.heat_map_tensor(point.view(-1, 2), h.view(-1, 1),
self.device, self.d,
self.heatmap_size)
heatmap = heat_map.cpu().detach().numpy().reshape(
(H_prime, W_prime))
start_H, end_H, start_W, end_W = utils.find_out_center_coords(
coords, self.patch_size, self.heatmap_size)
cell[start_H:end_H, start_W:end_W] += heatmap
idx = np.argwhere(heatmap != 0)
count[idx[:, 0] + start_H, idx[:, 1] + start_W] += 1
count[count == 0] = 1
cell = np.divide(cell, count)
return cell
def training_validation(self):
since = time.time()
for epoch in range(self.SC_CNN.start_epoch, self.epochs):
print(f"Epoch {epoch}/{self.epochs-1}\n", '-' * 10, "\n")
running_loss = 0.0
for phase in ['Train', 'Validation']:
self.SC_CNN.model.train(
True) if phase == 'Train' else self.SC_CNN.model.train(
False)
# Iterate over data.
for batch_idx, (inputs, heatmaps, epsilons) in enumerate(
self.dataLoaders[phase]):
start_batch = time.time()
inputs, heatmaps, epsilons = inputs.to(
self.device), heatmaps.to(self.device), epsilons.to(
self.device)
# forward
points, h_value = self.SC_CNN.model(inputs)
if self.auto_grad == 'PyTorch':
predicted = utils.heat_map_tensor(
points, h_value, self.device, self.d,
self.heatmap_size)
else:
# Amirali's implementation
predicted = self.SC_CNN.Map(points, h_value,
self.device, self.d,
self.heatmap_size,
self.version)
predicted = predicted.view(predicted.shape[0], -1)
heatmaps = heatmaps.view(heatmaps.shape[0], -1)
feature = heatmaps.shape[1]
epsilons = epsilons.unsqueeze(1)
epsilons = torch.repeat_interleave(epsilons,
repeats=feature,
dim=1)
weights = heatmaps + epsilons
if phase == 'Train':
loss = self.SC_CNN.criterion(predicted, heatmaps)
weight_loss = loss * weights
# Sum over one data
# Average over different data
loss = torch.sum(weight_loss, dim=1)
loss = torch.mean(loss)
running_loss += loss.item() * inputs.shape[0]
# Update parameters
self.SC_CNN.optimizer.zero_grad()
loss.backward()
self.SC_CNN.optimizer.step()
else:
loss = self.SC_CNN.val_criterion(predicted, heatmaps)
loss = torch.sqrt(loss)
loss = torch.sum(loss, dim=1)
loss = torch.sum(loss)
running_loss += loss.item()
end_batch = time.time()
if self.verbose:
print(
f"{phase}: epoch {epoch}: batch {batch_idx+1}/{int(np.ceil(self.dataset_sizes[phase]/self.batch_size))}:"
f"{end_batch-start_batch:3f} s {self.batch_size/(end_batch-start_batch):3f} data/s"
f"obj: {running_loss/((batch_idx+1)*self.batch_size):3f}-> Loss: {loss.item():3f}",
flush=True)
print("RAM: ", psutil.virtual_memory().percent)
# Empty Catch
torch.cuda.empty_cache()
epoch_loss = running_loss / self.dataset_sizes[phase]
print(f"{phase} -> Loss: {epoch_loss}", flush=True)
self.writer.add_scalars(f"{self.save}/loss",
{f"{phase}": epoch_loss}, epoch)
self.writer.flush()
if phase == 'Train':
self.SC_CNN.scheduler.step()
if phase == 'Validation' and epoch_loss < self.SC_CNN.best_loss:
self.SC_CNN.best_loss = epoch_loss
# Save the best model
utils.save_model(epoch, self.SC_CNN.model,
self.SC_CNN.optimizer,
self.SC_CNN.scheduler, epoch_loss,
self.save)
self.writer.close()
time_elapsed = time.time() - since
self.logger.info(
f"Training complete in {time_elapsed // 60:.0f}m {time_elapsed % 60:.0f}s"
)
self.logger.info(f"Best loss: {self.SC_CNN.best_loss:4f}")