def vis_on_batch(self, batch, savedir_image):
image = batch['images']
gt = np.asarray(batch['masks'], np.float32)
gt /= (gt.max() + 1e-8)
res = self.predict_on_batch(batch)
image = F.interpolate(image, size=gt.shape[-2:], mode='bilinear', align_corners=False)
img_res = hu.save_image(savedir_image,
hu.denormalize(image, mode='rgb')[0],
mask=res[0], return_image=True)
img_gt = hu.save_image(savedir_image,
hu.denormalize(image, mode='rgb')[0],
mask=gt[0], return_image=True)
img_gt = models.text_on_image( 'Groundtruth', np.array(img_gt), color=(0,0,0))
img_res = models.text_on_image( 'Prediction', np.array(img_res), color=(0,0,0))
if 'points' in batch:
pts = batch['points'][0].numpy().copy()
pts[pts == 1] = 2
pts[pts == 0] = 1
pts[pts == 255] = 0
img_gt = np.array(hu.save_image(savedir_image, img_gt/255.,
points=pts, radius=2, return_image=True))
img_list = [np.array(img_gt), np.array(img_res)]
hu.save_image(savedir_image, np.hstack(img_list))
def vis_on_batch(self, batch, savedir_image):
self.eval()
images = batch["images"].to(self.device)
points = batch["points"].long().to(self.device)
logits = self.model_base.forward(images)
probs = logits.sigmoid().cpu().numpy()
blobs = lcfcn_loss.get_blobs(probs=probs)
pred_counts = (np.unique(blobs)!=0).sum()
pred_blobs = blobs
pred_probs = probs.squeeze()
# loc
pred_count = pred_counts.ravel()[0]
pred_blobs = pred_blobs.squeeze()
pred_points = lcfcn_loss.blobs2points(pred_blobs).squeeze()
img_org = hu.get_image(batch["images"],denorm="rgb")
i1 = convert(np.array(img_org), batch['points'][0], enlarge=20)
i2 = convert(np.array(img_org), pred_blobs, enlarge=0)
i3 = convert(np.array(img_org), pred_points, enlarge=20)
hu.save_image(savedir_image, np.hstack([i1, i2, i3]))
def vis_on_loader(self, loader, savedir_images, epoch=None):
self.model.eval()
count = 0
for batch in tqdm.tqdm(loader):
if batch['label'][0][1:].sum() == 0:
continue
cam_dict = self.make_multiscale_cam(batch)
keys = cam_dict['keys']
high_res = cam_dict['high_res']
im_rgb = batch['original'][0]
pred_cam = hu.f2l(hi.gray2cmap(cam_dict['cam_crf']))
hu.save_image(
os.path.join(
savedir_images,
batch['meta'][0]['name'] + '_cam_epoch:%d.jpg' % (epoch)),
0.5 * (np.array(im_rgb / 255.)) + 0.5 * pred_cam)
for i, k in enumerate(keys):
pred = hu.f2l(hi.gray2cmap(high_res[i]))
# im_rgb = hu.denormalize(im, 'rgb')
# im_rgb = hu.f2l(im_rgb)
hu.save_image(
os.path.join(
savedir_images, batch['meta'][0]['name'] +
'_class:%d_epoch:%d.jpg' % (k, epoch)),
0.5 * (np.array(im_rgb / 255.)) + 0.5 * pred)
if count > 3:
break
count += 1
def vis_on_batch(self, split, batch, savedir_images, epoch):
self.eval()
images = batch['images'].to(self.device)
labels = batch['labels'].to(self.device)
mean = torch.tensor(self.netC.dataset.mean).to(self.device)
std = torch.tensor(self.netC.dataset.std).to(self.device)
org = images.detach()
org = (org * std.view(1, 3, 1, 1) +
mean.view(1, 3, 1, 1)).cpu().numpy()
img_list = []
if split == 'train' and self.model_dict.get('netA') is not None:
augimages, _, _ = self.netA.apply_augmentation(images, labels)
aug = augimages.detach()
aug = (aug * std.view(1, 3, 1, 1) +
mean.view(1, 3, 1, 1)).cpu().numpy()
for i in range(org.shape[0]):
both = np.concatenate([org[i], aug[i]], axis=1)
img_list += [both]
img_list = np.concatenate(img_list, axis=2)
else:
img_list = np.concatenate(org, axis=2)
hu.save_image(
os.path.join(savedir_images, '%s_%s.jpg' % (split, epoch)),
img_list)
def save(self, img_pil, y_list, x_list, mask):
img_p = hi.points_on_image(y_list, x_list, img_pil)
img_maskspil = hi.mask_on_image(img_p,
mask.astype('uint8'),
add_bbox=True)
hu.save_image('masker.jpg', img_maskspil)
def save_tmp(fname, images, logits, points):
from haven import haven_utils as hu
probs = F.softmax(logits, 1)
mask = probs.argmax(dim=1).cpu().numpy().astype('uint8').squeeze() * 255
img_mask = hu.save_image('tmp2.png',
hu.denormalize(images, mode='rgb'),
mask=mask,
return_image=True)
hu.save_image(fname, np.array(img_mask) / 255., points=att_dict['points'])
def vis_on_batch(self, batch, savedir_image):
self.eval()
images = batch["images"].to('cpu') #helen changed this from .cuda()
#print(images)
#print(images.shape)
points = batch["points"].long().to(
'cpu') #helen changed this from .cuda()
logits = self.model.forward(images)
probs = logits.sigmoid().cpu().numpy()
blobs = lcfcn_loss.get_blobs(probs=probs)
pred_counts = (np.unique(blobs) != 0).sum()
pred_blobs = blobs
pred_probs = probs.squeeze()
# loc
pred_count = pred_counts.ravel()[0]
pred_blobs = pred_blobs.squeeze()
img_org = hu.get_image(batch["images"], denorm="rgb")
# true points
y_list, x_list = np.where(batch["points"][0].long().numpy().squeeze())
img_peaks = hi.points_on_image(y_list, x_list, img_org, radius=11)
text = "%s ground truth" % (batch["points"].sum().item())
hi.text_on_image(text=text, image=img_peaks)
# pred points
pred_points = lcfcn_loss.blobs2points(pred_blobs).squeeze()
y_list, x_list = np.where(pred_points.squeeze())
img_pred = hi.mask_on_image(img_org, pred_blobs)
#img_pred = hi.points_on_image(y_list, x_list, img_org)
text = "%s predicted" % (len(y_list))
hi.text_on_image(text=text, image=img_pred)
# *************** helen added this code
plt.imshow(
img_pred
) #these lines of code display the image with the model's predications on it
plt.show()
# *************** end of helen's code
# heatmap
heatmap = hi.gray2cmap(pred_probs)
heatmap = hu.f2l(heatmap)
hi.text_on_image(text="lcfcn heatmap", image=heatmap)
img_mask = np.hstack([img_peaks, img_pred, heatmap])
hu.save_image(savedir_image, img_mask)
def vis_on_batch_helen(self, batch, savedir_image):
self.eval()
images = batch.to('cpu')
#print(images) #print statement to check what images is for debugging
#print(images.shape) #print statement to check images has the correct shape
logits = self.model.forward(images)
probs = logits.sigmoid().cpu().numpy()
blobs = lcfcn_loss.get_blobs(probs=probs)
pred_counts = (np.unique(blobs) != 0).sum()
pred_blobs = blobs
pred_probs = probs.squeeze()
# loc
pred_count = pred_counts.ravel()[0]
pred_blobs = pred_blobs.squeeze()
img_org = hu.get_image(images, denorm="rgb")
#this is what was originally written: img_org = hu.get_image(batch["images"],denorm="rgb")
#the following lines are commented out because my own data does not have labels and will thus throw errors
# true points
#y_list, x_list = np.where(batch["points"][0].long().numpy().squeeze())
#img_peaks = hi.points_on_image(y_list, x_list, img_org, radius=11)
#text = "%s ground truth" % (batch["points"].sum().item())
#hi.text_on_image(text=text, image=img_peaks)
# pred points
pred_points = lcfcn_loss.blobs2points(pred_blobs).squeeze()
y_list, x_list = np.where(pred_points.squeeze())
img_pred = hi.mask_on_image(img_org, pred_blobs)
# img_pred = hi.points_on_image(y_list, x_list, img_org)
text = "%s predicted" % (len(y_list))
hi.text_on_image(text=text, image=img_pred)
# these lines of code display the image with the model's predications on it
plt.imshow(img_pred)
plt.show()
# heatmap
heatmap = hi.gray2cmap(pred_probs)
heatmap = hu.f2l(heatmap)
hi.text_on_image(text="lcfcn heatmap", image=heatmap)
img_mask = np.hstack([img_pred, heatmap]) #helen took out im_peaks
#this is what was originally written: img_mask = np.hstack([img_peaks, img_pred, heatmap])
hu.save_image(savedir_image, img_mask)
def get_blob_list(mask_dict, points_mask, img_pil, split_inst=False):
blob_list = []
mask = preds = mask_dict['preds']
probs = mask_dict['probs']
assert probs.shape[1] == preds.shape[0]
assert probs.shape[2] == preds.shape[1]
imask = np.zeros(mask.shape)
cmask = np.zeros(mask.shape)
blob_id = 1
for c in np.unique(mask):
if c == 0:
continue
probs_class = probs[c]
point_ind = points_mask == c
inst_mask = morphology.label(mask==c)
for l in np.unique(inst_mask):
if l == 0:
continue
blob_ind = inst_mask == l
locs = np.where(blob_ind * point_ind)
y_list, x_list = locs
n_points = len(y_list)
if n_points == 0:
continue
if n_points > 1 and split_inst:
# split multiple points
img_points = hi.points_on_image(y_list, x_list, img_pil)
img_masks = hi.mask_on_image(img_pil, mask)
img_masks = hi.mask_on_image(img_points.copy(), blob_ind)
hu.save_image('tmp.jpg', (img_points)*0.5 + hu.f2l(hi.gray2cmap(probs_class))*0.5)
hu.save_image('tmp.jpg', img_masks)
for yi, xi in zip(y_list, x_list):
imask, cmask, blob_list, blob_id = add_mask(yi, xi, points_mask,
blob_ind,
n_points, blob_list, imask, cmask,
blob_id)
else:
# add for that single point
yi, xi = y_list[0], x_list[0]
imask, cmask, blob_list, blob_id = add_mask(yi, xi, points_mask,
blob_ind,
n_points, blob_list, imask, cmask,
blob_id)
return blob_list, cmask.astype('uint8'), imask.astype('uint8')
def vis_on_batch(self, batch, savedir_image):
image = batch['images']
res = self.predict_on_batch(batch)
# img_gt = hu.save_image('tmp.png',
# hu.denormalize(image, mode='rgb')[0],
# points=batch['points'][0], radius=1,
# return_image=True)
img_res = hu.save_image(savedir_image,
hu.denormalize(image, mode='rgb')[0],
mask=res.cpu().numpy(), return_image=True)
# img_gt = semseg.text_on_image( 'Groundtruth', np.array(img_gt), color=(0,0,0))
# img_res = semseg.text_on_image( 'Prediction', np.array(img_res), color=(0,0,0))
hu.save_image(savedir_image, np.hstack([np.array(img_res)]))
def vis_on_batch(self, batch, savedir_image):
from skimage.segmentation import mark_boundaries
from skimage import data, io, segmentation, color
from skimage.measure import label
self.eval()
pred_mask = self.predict_on_batch(batch)
img = hu.get_image(batch["images"], denorm="rgb")
img_np = np.array(img)
pm = pred_mask.squeeze()
out = color.label2rgb(label(pm), image=(img_np), image_alpha=1.0, bg_label=0)
img_mask = mark_boundaries(out.squeeze(), label(pm).squeeze())
out = color.label2rgb(label(batch["mask_classes"][0]), image=(img_np), image_alpha=1.0, bg_label=0)
img_gt = mark_boundaries(out.squeeze(), label(batch["mask_classes"]).squeeze())
hu.save_image(savedir_image, np.hstack([img_gt, img_mask]))
def save_example_results(savedir_base="results"):
import os
import pandas
import requests
import io
import matplotlib.pyplot as plt
from .. import haven_results as hr
from .. import haven_utils as hu
from PIL import Image
# create hyperparameters
exp_list = [{
"dataset": "mnist",
"model": "mlp",
"lr": lr
} for lr in [1e-1, 1e-2, 1e-3]]
for i, exp_dict in enumerate(exp_list):
# get hash for experiment
exp_id = hu.hash_dict(exp_dict)
# add scores for loss, and accuracy
score_list = []
for e in range(1, 10):
score_list += [{
"epoch": e,
"loss": 1 - e * exp_dict["lr"] * 0.9,
"acc": e * exp_dict["lr"] * 0.1
}]
# save scores and images
hu.save_json(os.path.join(savedir_base, exp_id, "exp_dict.json"),
exp_dict)
hu.save_pkl(os.path.join(savedir_base, exp_id, "score_list.pkl"),
score_list)
url = "https://raw.githubusercontent.com/haven-ai/haven-ai/master/haven/haven_examples/data/%d.png" % (
i + 1)
response = requests.get(url).content
img = plt.imread(io.BytesIO(response), format="JPG")
hu.save_image(os.path.join(savedir_base, exp_id, "images/1.png"),
img[:, :, :3])
def vis_on_batch(self, batch, savedir_image):
self.eval()
images = batch["images"].cuda()
points = batch["points"].long().cuda()
logits = self.model_base.forward(images)
probs = logits.sigmoid().cpu().numpy()
blobs = lcfcn_loss.get_blobs(probs=probs)
pred_counts = (np.unique(blobs)!=0).sum()
pred_blobs = blobs
pred_probs = probs.squeeze()
# loc
pred_count = pred_counts.ravel()[0]
pred_blobs = pred_blobs.squeeze()
img_org = hu.get_image(batch["images"],denorm="rgb")
# true points
y_list, x_list = np.where(batch["points"][0].long().numpy().squeeze())
img_peaks = haven_img.points_on_image(y_list, x_list, img_org)
text = "%s ground truth" % (batch["points"].sum().item())
haven_img.text_on_image(text=text, image=img_peaks)
# pred points
pred_points = lcfcn_loss.blobs2points(pred_blobs).squeeze()
y_list, x_list = np.where(pred_points.squeeze())
img_pred = hi.mask_on_image(img_org, pred_blobs)
# img_pred = haven_img.points_on_image(y_list, x_list, img_org)
text = "%s predicted" % (len(y_list))
haven_img.text_on_image(text=text, image=img_pred)
# heatmap
heatmap = hi.gray2cmap(pred_probs)
heatmap = hu.f2l(heatmap)
haven_img.text_on_image(text="lcfcn heatmap", image=heatmap)
img_mask = np.hstack([img_peaks, img_pred, heatmap])
hu.save_image(savedir_image, img_mask)
def vis_on_batch(self, batch, savedir_image, save_preds=False):
self.eval()
images = batch["images"].cuda()
counts = float(batch["counts"][0])
logits = self.model_base(images)
probs = logits.sigmoid()
# get points from attention
att_dict = self.att_model.get_attention_dict(
images_original=torch.FloatTensor(
hu.denormalize(batch['images'], mode='rgb')),
counts=batch['counts'][0],
probs=probs.squeeze(),
return_roi=True)
blobs = lcfcn.get_blobs(probs.squeeze().detach().cpu().numpy())
org_img = hu.denormalize(images.squeeze(), mode='rgb')
rgb_labels = label2rgb(
blobs,
hu.f2l(org_img),
bg_label=0,
bg_color=None,
)
res1 = mark_boundaries(rgb_labels, blobs)
if att_dict['roi_mask'] is not None:
img_mask = hu.save_image('tmp2.png',
org_img,
mask=att_dict['roi_mask'] == 1,
return_image=True)
res2 = hu.save_image('tmp.png',
np.array(img_mask) / 255.,
points=att_dict['points'],
radius=1,
return_image=True)
os.makedirs(os.path.dirname(savedir_image), exist_ok=True)
# plt.savefig(savedir_image.replace('.jpg', '.png')
hu.save_image(savedir_image.replace('.jpg', '.png'),
np.hstack([res1, np.array(res2) / 255.]))
def save_example_results(savedir_base='results'):
import os, pandas
import requests, io
import matplotlib.pyplot as plt
from .. import haven_results as hr
from .. import haven_utils as hu
from PIL import Image
# create hyperparameters
exp_list = [{
'dataset': 'mnist',
'model': 'mlp',
'lr': lr
} for lr in [1e-1, 1e-2, 1e-3]]
for i, exp_dict in enumerate(exp_list):
# get hash for experiment
exp_id = hu.hash_dict(exp_dict)
# add scores for loss, and accuracy
score_list = []
for e in range(1, 10):
score_list += [{
'epoch': e,
'loss': 1 - e * exp_dict['lr'] * 0.9,
'acc': e * exp_dict['lr'] * 0.1
}]
# save scores and images
hu.save_json(os.path.join(savedir_base, exp_id, 'exp_dict.json'),
exp_dict)
hu.save_pkl(os.path.join(savedir_base, exp_id, 'score_list.pkl'),
score_list)
url = 'https://raw.githubusercontent.com/haven-ai/haven-ai/master/haven/haven_examples/data/%d.png' % (
i + 1)
response = requests.get(url).content
img = plt.imread(io.BytesIO(response), format='JPG')
hu.save_image(os.path.join(savedir_base, exp_id, 'images/1.png'),
img[:, :, :3])
def vis_on_batch(self, batch, savedir):
self.eval()
os.makedirs(savedir, exist_ok=True)
overlayed = hi.mask_on_image(batch['image_pil'][0],
np.array(batch['inst_pil'][0]),
add_bbox=True)
overlayed = Image.fromarray((overlayed * 255).astype('uint8'))
images = batch['images']
img = images[0]
prediction = self.model([img.to('cuda:0')])
org_img = Image.fromarray(img.mul(255).permute(1, 2, 0).byte().numpy())
pred_list = []
score_list = prediction[0]['scores']
for i in range(len(score_list)):
if score_list[i] < 0.5:
break
pred = ((prediction[0]['masks'][i, 0] >
0.5).mul(255)).byte().cpu().numpy()
pred_list += [Image.fromarray(pred)]
img_name = batch['meta'][0]['name']
for pred in pred_list:
org_img = Image.fromarray(
(hi.mask_on_image(org_img, pred) * 255).astype('uint8'))
fname = os.path.join(savedir, '%s.jpg' % img_name)
overlayed = hi.text_on_image(
'gt', np.array(overlayed.resize((350, 200), 2).copy()))
org_img = hi.text_on_image('preds',
np.array(org_img.resize((350, 200))).copy())
img = np.concatenate([org_img, overlayed.astype(float)],
axis=1).astype('float32') / 255.
hu.save_image(fname=fname, img=img)
print('image saved: %s' % fname)
def vis_on_batch(self, batch, savedir_image):
image = batch['images']
original = hu.denormalize(image, mode='rgb')[0]
img_pred = hu.save_image(savedir_image,
original,
mask=self.predict_on_batch(batch, method='semseg'), return_image=True)
img_gt = hu.save_image(savedir_image,
original,
return_image=True)
gt_counts = float(batch['points'].sum())
pred_counts = self.predict_on_batch(batch, method='counts')
img_gt = models.text_on_image( 'Groundtruth: %d' % gt_counts, np.array(img_gt), color=(0,0,0))
img_pred = models.text_on_image( 'Prediction: %d' % pred_counts, np.array(img_pred), color=(0,0,0))
if 'points' in batch:
pts = (batch['points'][0].numpy().copy() != 0).astype('uint8')
img_gt = np.array(hu.save_image(savedir_image, img_gt/255.,
points=pts.squeeze(), radius=2, return_image=True))
img_list = [np.array(img_gt), np.array(img_pred)]
hu.save_image("%s" %(savedir_image), np.hstack(img_list))
def vis_on_loader(self, loader, **extras):
self.eval()
for batch in loader:
images, labels = batch
probs = torch.softmax(self.model.forward(
images.view(images.shape[0], -1)),
dim=1)
score, label = probs.max(dim=1)
i_list = []
for i in range(probs.shape[0]):
pil_img = hu.save_image("tmp", images[i], return_image=True)
img = get_image(
pil_img,
"Predicted %d (Prob: %.2f)" % (label[i], score[i]))
i_list += [img]
if i > 5:
break
return np.hstack(i_list)[:, :, None].repeat(3, axis=2)
def save_images(exp_dict):
dataset_name = exp_dict['dataset']['name']
n_classes = exp_dict['dataset']['n_classes']
model = models.get_model(model_dict=exp_dict['model'],
exp_dict=exp_dict,
train_set=None).cuda()
state_dict = hc.load_checkpoint(exp_dict,
savedir_base,
fname='model_best.pth')
model.load_state_dict(state_dict)
model.eval()
np.random.seed(1)
train_set = datasets.get_dataset(dataset_dict={'name': dataset_name},
datadir=None,
split="test",
exp_dict=exp_dict)
n_images = 0
for _ in range(len(train_set)):
i = np.random.choice(len(train_set))
b = train_set[i]
if n_images > 5:
break
if b['masks'].sum() == 0:
print(i)
continue
n_images += 1
batch = ut.collate_fn([b])
image = batch['images']
gt = np.asarray(batch['masks'], np.float32)
gt /= (gt.max() + 1e-8)
image = F.interpolate(image,
size=gt.shape[-2:],
mode='bilinear',
align_corners=False)
img_rgb = hu.f2l(hu.denormalize(image, mode='rgb')[0])
img_rgb = (np.array(img_rgb) * 255.).astype('uint8')
# save rgb
fname_rgb = '.tmp/covid_qualitative/%s/%s/%d_rgb.png' % (exp_group,
'gt', i)
hu.save_image(fname_rgb, img_rgb)
# save pts
fname_pts = '.tmp/covid_qualitative/%s/%s/%d_pts.png' % (exp_group,
'gt', i)
img_gt = np.array(hu.save_image('', img_rgb, return_image=True))
if 'points' in batch:
pts = batch['points'][0].numpy()
pts[pts == 1] = 2
pts[pts == 0] = 1
pts[pts == 255] = 0
img_gt = np.array(
hu.save_image('',
img_gt / 255.,
points=pts,
radius=2,
return_image=True))
hu.save_image(fname_pts, img_gt)
# save mask
fname_mask = '.tmp/covid_qualitative/%s/%s/%d_mask.png' % (exp_group,
'gt', i)
img_mask = np.array(
hu.save_image('', img_rgb, mask=gt[0], return_image=True))
hu.save_image(fname_mask, img_mask)
# pred
fname_pred = '.tmp/covid_qualitative/%s/%s/%d_%s.png' % (
exp_group, 'preds', i, exp_dict['model']['loss'])
res = model.predict_on_batch(batch)
img_res = hu.save_image('', img_rgb, mask=res[0], return_image=True)
hu.save_image(fname_pred, np.array(img_res))
batch = ut.collate_fn([b])
image = batch['images']
gt = np.asarray(batch['masks'], np.float32)
gt /= (gt.max() + 1e-8)
image = F.interpolate(image,
size=gt.shape[-2:],
mode='bilinear',
align_corners=False)
# img_res = hu.save_image('',
# hu.denormalize(image, mode='rgb')[0],
# mask=res[0], return_image=True)
img_gt = hu.save_image('',
hu.denormalize(image, mode='rgb')[0],
mask=gt[0],
return_image=True)
img_gt = models.text_on_image('Groundtruth',
np.array(img_gt),
color=(0, 0, 0))
# img_res = models.text_on_image( 'Prediction', np.array(img_res), color=(0,0,0))
if 'points' in batch:
img_gt = np.array(
hu.save_image('',
img_gt / 255.,
points=batch['points'][0].numpy() != 255,
radius=2,
return_image=True))
img_list = [np.array(img_gt)]
hu.save_image('.tmp/covid_datasets/%s.png' % exp_group,
test_loader = DataLoader(
test_set,
# sampler=val_sampler,
batch_size=1,
collate_fn=ut.collate_fn,
num_workers=0)
for i, batch in enumerate(test_loader):
points = (batch['points'].squeeze() == 1).numpy()
if points.sum() == 0:
continue
savedir_image = os.path.join('.tmp/qualitative/%d.png' % (i))
img = hu.denormalize(batch['images'], mode='rgb')
img_org = np.array(
hu.save_image(savedir_image,
img,
mask=batch['masks'].numpy(),
return_image=True))
img_list = [img_org]
with torch.no_grad():
for hash_id in hash_list:
score_path = os.path.join(savedir_base, hash_id,
'score_list_best.pkl')
score_list = hu.load_pkl(score_path)
exp_dict = hu.load_json(
os.path.join(savedir_base, hash_id, 'exp_dict.json'))
print(i, exp_dict['model']['loss'],
exp_dict['model'].get('with_affinity'), 'score:',
score_list[-1]['test_class1'])
# loop over the val_loader and saves image
for i, batch in enumerate(test_loader):
savedir_image = os.path.join("%s" % savedir, "save_preds",
"%s" % hash_dir, "%s" % split,
"%d.png" % i)
image = batch['images']
original = hu.denormalize(image, mode='rgb')[0]
gt = np.asarray(batch['masks'])
image = F.interpolate(image,
size=gt.shape[-2:],
mode='bilinear',
align_corners=False)
img_pred = hu.save_image(savedir_image,
original,
mask=model.predict_on_batch(batch),
return_image=True)
img_gt = hu.save_image(savedir_image,
original,
mask=gt,
return_image=True)
# add text_on_image here
img_gt = models.text_on_image('', np.array(img_gt), color=(0, 0, 0))
img_pred = models.text_on_image('',
np.array(img_pred),
color=(0, 0, 0))
if 'points' in batch:
pts = (batch['points'][0].numpy().copy()).astype('uint8')
def vis_on_batch(self, batch, savedir_image):
image = batch['images']
index = batch['meta'][0]['index']
gt = np.asarray(batch['masks'], np.float32)
gt /= (gt.max() + 1e-8)
res = self.predict_on_batch(batch)
image = F.interpolate(image,
size=gt.shape[-2:],
mode='bilinear',
align_corners=False)
original = hu.denormalize(image, mode='rgb')[0]
img_res = hu.save_image(savedir_image,
original,
mask=res[0],
return_image=True)
img_gt = hu.save_image(savedir_image,
original,
mask=gt[0],
return_image=True)
img_gt = models.text_on_image('Groundtruth',
np.array(img_gt),
color=(0, 0, 0))
img_res = models.text_on_image('Prediction',
np.array(img_res),
color=(0, 0, 0))
if 'points' in batch:
pts = batch['points'][0].numpy().copy()
pts[pts == 1] = 2
pts[pts == 0] = 1
pts[pts == 255] = 0
img_gt = np.array(
hu.save_image(savedir_image,
img_gt / 255.,
points=pts,
radius=2,
return_image=True))
# score map
if self.heuristic != 'random':
score_map = self.compute_uncertainty(batch['images'],
replicate=True,
scale_factor=1,
method=self.heuristic)
score_map = F.interpolate(score_map[None],
size=gt.shape[-2:],
mode='bilinear',
align_corners=False).squeeze()
h, w = score_map.shape
bbox_yxyx = get_rect_bbox(h, w, n_regions=self.n_regions)
heatmap = hu.f2l(
hi.gray2cmap(
(score_map / score_map.max()).cpu().numpy().squeeze()))
s_list = np.zeros(len(bbox_yxyx))
for i, (y1, x1, y2, x2) in enumerate(bbox_yxyx):
s_list[i] = score_map[y1:y2, x1:x2].mean()
img_bbox = bbox_yxyx_on_image(bbox_yxyx[[s_list.argmax()]],
original)
img_score_map = img_bbox * 0.5 + heatmap * 0.5
img_list = [
np.array(img_gt),
np.array(img_res), (img_score_map * 255).astype('uint8')
]
else:
img_list = [np.array(img_gt), np.array(img_res)]
hu.save_image(savedir_image, np.hstack(img_list))
# stop
model = models.get_model(model_dict=exp_dict['model'],
exp_dict=exp_dict,
train_set=train_set).cuda()
exp_id = hu.hash_dict(exp_dict)
fname = os.path.join('/mnt/public/results/toolkit/weak_supervision',
exp_id, 'model.pth')
model.model_base.load_state_dict(torch.load(fname)['model'], strict=False)
for k in range(5):
batch_id = np.where(train_set.labels)[0][k]
batch = ut.collate_fn([train_set[batch_id]])
logits = F.softmax(model.model_base.forward(batch['images'].cuda()),
dim=1)
img = batch['images'].cuda()
logits_new = model_aff.apply_affinity(batch['images'], logits, crf=0)
i1 = hu.save_image('old.png',
img=hu.denormalize(img, mode='rgb'),
mask=logits.argmax(dim=1).cpu().numpy(),
return_image=True)
i2 = hu.save_image('new.png',
img=hu.denormalize(img, mode='rgb'),
mask=logits_new.argmax(dim=1).cpu().numpy(),
return_image=True)
hu.save_image('tmp/tmp%d.png' % k,
np.concatenate([np.array(i1), np.array(i2)], axis=1))
print('saved %d' % k)
def train_on_batch(self, batch, **extras):
self.train()
images = batch["images"].cuda()
counts = float(batch["counts"][0])
logits = self.model_base(images)
if self.exp_dict['model'].get('loss') == 'lcfcn':
loss = lcfcn.compute_lcfcn_loss(logits, batch["points"].cuda(),
None)
probs = F.softmax(logits, 1)
mask = probs.argmax(
dim=1).cpu().numpy().astype('uint8').squeeze() * 255
# img_mask=hu.save_image('tmp2.png',
# hu.denormalize(images, mode='rgb'), mask=mask, return_image=True)
# hu.save_image('tmp2.png',np.array(img_mask)/255. , radius=3,
# points=batch["points"])
elif self.exp_dict['model'].get('loss') == 'glance':
pred_counts = logits[:, 1].mean()
loss = F.mse_loss(pred_counts.squeeze(), counts.float().squeeze())
elif self.exp_dict['model'].get('loss') == 'att_lcfcn':
probs = logits.sigmoid()
# get points from attention
att_dict = self.att_model.get_attention_dict(
images_original=torch.FloatTensor(
hu.denormalize(batch['images'], mode='rgb')),
counts=batch['counts'][0],
probs=probs.squeeze(),
return_roi=True)
if 1:
blobs = lcfcn.get_blobs(probs.squeeze().detach().cpu().numpy())
org_img = hu.denormalize(images.squeeze(), mode='rgb')
rgb_labels = label2rgb(
blobs,
hu.f2l(org_img),
bg_label=0,
bg_color=None,
)
res1 = mark_boundaries(rgb_labels, blobs)
img_mask = hu.save_image('tmp2.png',
org_img,
return_image=True)
res2 = hu.save_image('tmp.png',
np.array(img_mask) / 255.,
points=att_dict['points'],
radius=1,
return_image=True)
hu.save_image('tmp_blobs.png',
np.hstack([res1, np.array(res2) / 255.]))
loss = lcfcn.compute_loss(
probs=probs,
# batch["points"].cuda(),
points=att_dict['points'].cuda(),
roi_mask=att_dict['roi_mask'])
# loss += .5 * F.cross_entropy(logits,
# torch.from_numpy(1 -
# att_dict['mask_bg']).long().cuda()[None],
# ignore_index=1)
self.opt.zero_grad()
loss.backward()
if self.exp_dict['optimizer'] == 'sps':
self.opt.step(loss=loss)
else:
self.opt.step()
return {"train_loss": float(loss)}
model = models.get_model(model_dict=exp_dict['model'],
exp_dict=exp_dict,
train_set=train_set).cuda()
model_path = os.path.join(savedir_base, hash_id, 'model_best.pth')
# load best model
model.load_state_dict(hu.torch_load(model_path))
# loop over the val_loader and saves image
# get counts
counts = []
for i, batch in enumerate(test_loader):
count = float((batch['points'] == 1).sum())
counts += [count]
hu.save_image('.tmp/counts/%d_%d.png' %
(i, len(batch['point_list'][0]) // 2),
hu.denormalize(batch['images'], mode='rgb'),
mask=batch['masks'].numpy())
counts = np.array(counts)
val_dict = {}
val_dict_lst = []
for c in np.unique(counts):
val_meter = metrics.SegMeter(split=test_loader.dataset.split)
for i, batch in enumerate(tqdm.tqdm(test_loader)):
count = float((batch['points'] == 1).sum())
if count != c:
continue
val_meter.val_on_batch(model, batch)
score_dict = val_meter.get_avg_score()
if __name__ == "__main__":
savedir_base = '/mnt/public/results/toolkit/weak_supervision'
hash_list = ['a55d2c5dda331b1a0e191b104406dd1c']
# LCFCN
# hash_id = 'bcba046296675e9e3af5cd9f353d217b'
for hash_id in hash_list:
exp_dict = hu.load_json(
os.path.join(savedir_base, hash_id, 'exp_dict.json'))
datadir = '/mnt/public/datasets/DeepFish/'
split = 'train'
train_set = datasets.get_dataset(dataset_dict=exp_dict["dataset"],
split=split,
datadir=datadir,
exp_dict=exp_dict,
dataset_size=exp_dict['dataset_size'])
train_loader = DataLoader(
train_set,
# sampler=val_sampler,
batch_size=1,
collate_fn=ut.collate_fn,
num_workers=0)
for i, batch in enumerate(train_loader):
points = (batch['points'].squeeze() == 1).numpy()
if points.sum() == 0:
continue
savedir_image = os.path.join('.tmp/habitats/%s/%d.png' %
(batch['meta'][0]['habitat'], i))
img = hu.denormalize(batch['images'], mode='rgb')
# img_pred = model.predict_on_batch(batch)
hu.save_image(savedir_image, img, points=points, radius=1)
