Spaces:
Runtime error
Runtime error
File size: 4,976 Bytes
7629b39 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 |
# code idea from https://github.com/sergeyprokudin/bps
import os
import numpy as np
from PIL import Image
import time
import scipy
import scipy.spatial
import pymp
#####################
QUERY_POINTS = np.asarray([30, 34, 31, 55, 29, 84, 35, 108, 34, 145, 29, 171, 27,
196, 29, 228, 58, 35, 61, 55, 57, 83, 56, 109, 63, 148, 58, 164, 57, 197, 60,
227, 81, 26, 87, 58, 85, 87, 89, 117, 86, 142, 89, 172, 84, 197, 88, 227, 113,
32, 116, 58, 112, 88, 118, 113, 109, 147, 114, 173, 119, 201, 113, 229, 139,
29, 141, 59, 142, 93, 139, 117, 146, 147, 141, 173, 142, 201, 143, 227, 170,
26, 173, 59, 166, 90, 174, 117, 176, 141, 169, 175, 167, 198, 172, 227, 198,
30, 195, 59, 204, 85, 198, 116, 195, 140, 198, 175, 194, 193, 199, 227, 221,
26, 223, 57, 227, 83, 227, 113, 227, 140, 226, 173, 230, 196, 228, 229]).reshape((64, 2))
#####################
class SegBPS():
def __init__(self, query_points=QUERY_POINTS, size=256):
self.size = size
self.query_points = query_points
row, col = np.indices((self.size, self.size))
self.indices_rc = np.stack((row, col), axis=2) # (256, 256, 2)
self.pts_aranged = np.arange(64)
return
def _do_kdtree(self, combined_x_y_arrays, points):
# see https://stackoverflow.com/questions/10818546/finding-index-of-nearest-
# point-in-numpy-arrays-of-x-and-y-coordinates
mytree = scipy.spatial.cKDTree(combined_x_y_arrays)
dist, indexes = mytree.query(points)
return indexes
def calculate_bps_points(self, seg, thr=0.5, vis=False, out_path=None):
# seg: input segmentation image of shape (256, 256) with values between 0 and 1
query_val = seg[self.query_points[:, 0], self.query_points[:, 1]]
pts_fg = self.pts_aranged[query_val>=thr]
pts_bg = self.pts_aranged[query_val<thr]
candidate_inds_bg = self.indices_rc[seg<thr]
candidate_inds_fg = self.indices_rc[seg>=thr]
if candidate_inds_bg.shape[0] == 0:
candidate_inds_bg = np.ones((1, 2)) * 128 # np.zeros((1, 2))
if candidate_inds_fg.shape[0] == 0:
candidate_inds_fg = np.ones((1, 2)) * 128 # np.zeros((1, 2))
# calculate nearest points
all_nearest_points = np.zeros((64, 2))
all_nearest_points[pts_fg, :] = candidate_inds_bg[self._do_kdtree(candidate_inds_bg, self.query_points[pts_fg, :]), :]
all_nearest_points[pts_bg, :] = candidate_inds_fg[self._do_kdtree(candidate_inds_fg, self.query_points[pts_bg, :]), :]
all_nearest_points_01 = all_nearest_points / 255.
if vis:
self.visualize_result(seg, all_nearest_points, out_path=out_path)
return all_nearest_points_01
def calculate_bps_points_batch(self, seg_batch, thr=0.5, vis=False, out_path=None):
# seg_batch: input segmentation image of shape (bs, 256, 256) with values between 0 and 1
bs = seg_batch.shape[0]
all_nearest_points_01_batch = np.zeros((bs, self.query_points.shape[0], 2))
for ind in range(0, bs): # 0.25
seg = seg_batch[ind, :, :]
all_nearest_points_01 = self.calculate_bps_points(seg, thr=thr, vis=vis, out_path=out_path)
all_nearest_points_01_batch[ind, :, :] = all_nearest_points_01
return all_nearest_points_01_batch
def visualize_result(self, seg, all_nearest_points, out_path=None):
import matplotlib as mpl
mpl.use('Agg')
import matplotlib.pyplot as plt
# img: (256, 256, 3)
img = (np.stack((seg, seg, seg), axis=2) * 155).astype(np.int)
if out_path is None:
ind_img = 0
out_path = '../test_img' + str(ind_img) + '.png'
fig, ax = plt.subplots()
plt.imshow(img)
plt.gca().set_axis_off()
plt.subplots_adjust(top = 1, bottom = 0, right = 1, left = 0, hspace = 0, wspace = 0)
plt.margins(0,0)
ratio_in_out = 1 # 255
for idx, (y, x) in enumerate(self.query_points):
x = int(x*ratio_in_out)
y = int(y*ratio_in_out)
plt.scatter([x], [y], marker="x", s=50)
x2 = int(all_nearest_points[idx, 1])
y2 = int(all_nearest_points[idx, 0])
plt.scatter([x2], [y2], marker="o", s=50)
plt.plot([x, x2], [y, y2])
plt.savefig(out_path, bbox_inches='tight', pad_inches=0)
plt.close()
return
if __name__ == "__main__":
ind_img = 2 # 4
path_seg_top = '...../pytorch-stacked-hourglass/results/dogs_hg8_ks_24_v1/test/'
path_seg = os.path.join(path_seg_top, 'seg_big_' + str(ind_img) + '.png')
img = np.asarray(Image.open(path_seg))
# min is 0.004, max is 0.9
# low values are background, high values are foreground
seg = img[:, :, 1] / 255.
# calculate points
bps = SegBPS()
bps.calculate_bps_points(seg, thr=0.5, vis=False, out_path=None)
|