def __init__(self, img=None, ann=None, **kwargs): if not img is None: self.img = img elif "npimg" in kwargs: self.img = Image.fromarray(kwargs["npimg"].astype('uint8')) else: self.img = Image.open(kwargs["img_path"]) if ann is not None: self.ann = ann elif "ann_path" in kwargs: scale = kwargs.get("scale", None) self.sc = scale #print("AnnImg scale", scale) self.ann = Ann(file=kwargs["ann_path"], scale=scale) else: self.ann = None
def ann(self, iid): ann_path = self.ann_path.format(iid=iid) scale = self.scales.get(iid, None) if not os.path.exists(ann_path): ann_path = None return None, iid else: return Ann(file=ann_path, scale=scale), iid
def fun(timg): img, tgt = timg.ssd() if cuda: img = img.cuda() dets = model.forward(Variable(img)) dim = img.size(2) dets = dets[0].data if cuda: dets = dets.cpu() ann = Ann(tensor=dets, dim=img.size(2)) return ann
def merge(self, fun, d=900, s=800, mimg=None): ret = [] for timg, x, y in self.mtile(mimg, d, s): ann = fun(timg) if ann is None: continue ann.dxy(-x, -y) ret.append(ann.dets) ret = np.concatenate(ret) return Ann(dets=ret)
def posmid(self, d, dist=None): if dist is None: dist = 1.5 * d pos = list([aimg.np() for cl, aimg in self.crop_ann(d)]) dets, _ = self.ann.midd(dist) ann = Ann(dets=dets) print("posmid", len(pos), ann.count) aimg = AnnImg(self.img, ann) mid = list([aimg.np() for cl, aimg in aimg.crop_ann(d)]) Xp, Xn = np.stack(pos), np.stack(mid) Yp, Yn = np.ones(Xp.shape[0]), np.zeros(Xn.shape[0]) return np.concatenate([Xp, Xn]), np.concatenate([Yp, Yn]), Xp, Xn
def neg(self, ratio=1.0, num=None, minn=None): if num is None: num = int(self.count * ratio) if minn is not None: num = max(num, minn) mask = self.neg_mask().inv().np() #print("mask", mask.shape) y, x = mask[0].nonzero() pts = np.stack([x, y], axis=1) np.random.shuffle(pts) xy = pts[:num] cl = -1 * np.ones((xy.shape[0], 1)) ann = Ann(dets=np.concatenate([cl, xy], axis=1).astype('int32')) return self.wann(ann)
def nimg(self, d, ratio=1): numneg = int(self.count * ratio) W, H = self.WH arr = [] steps = 0 while len(arr) < numneg and steps < 10000: steps += 1 x, y = (randint(0, W), randint(0, H)) if self.cropdd(x, y, d).count == 0: arr.append((-1, x, y)) ann = Ann(dets=np.array(arr)) return AnnImg(self.img, ann)
def hneg(self, ratio=1.0, num=None, minn=None, th=False, conf=10): if num is None: num = int(self.count * ratio) if minn is not None: num = max(num, minn) mask = self.neg_mask().inv().np() #print("mask", mask.shape) y, x = mask[0].nonzero() hm = self.np()[0] #print("hm", hm.shape) idx = np.argsort(hm[y, x]) py, px = y[idx], x[idx] xy = np.stack([px, py], axis=1) cl = -1 * np.ones((xy.shape[0], 1)) h = hm[xy[:, 1], xy[:, 0]] ann = Ann(dets=np.concatenate([cl.reshape(-1, 1), xy], axis=1).astype( 'int32')) ann.conf = ((255 - h) * 100) / 255 ann = ann.fconf(conf) #print(ann) if ann.count < num: return self.wann(ann) if th: ann = ann.pnms(50) return self.wann(ann.take(num))
def open(cls, img_path, ann_path=None, scale=None): ann = None if not ann_path is None: ann = Ann(file=ann_path, scale=scale) return AnnImg(img, ann)
class AnnImg(object): def __init__(self, img=None, ann=None, **kwargs): if not img is None: self.img = img elif "npimg" in kwargs: self.img = Image.fromarray(kwargs["npimg"].astype('uint8')) else: self.img = Image.open(kwargs["img_path"]) if ann is not None: self.ann = ann elif "ann_path" in kwargs: scale = kwargs.get("scale", None) self.sc = scale #print("AnnImg scale", scale) self.ann = Ann(file=kwargs["ann_path"], scale=scale) else: self.ann = None @classmethod def open(cls, img_path, ann_path=None, scale=None): ann = None if not ann_path is None: ann = Ann(file=ann_path, scale=scale) return AnnImg(img, ann) def crop(self, x1, y1, x2, y2, truncate=False, overlap=True): img = self.img.crop((x1, y1, x2, y2)) ann = None if self.ann is not None: ann = self.ann.crop(x1, y1, x2, y2, truncate=truncate, overlap=overlap) return AnnImg(img, ann) def cropH(self): W, H = self.WH img1 = self.crop(0, 0, W, H // 2) img2 = self.crop(0, H // 2, W, H) return img1, img2 def cropW(self): W, H = self.WH img1 = self.crop(0, 0, W // 2, H) img2 = self.crop(W // 2, 0, W, H) return img1, img2 def cropHW(self): W, H = self.WH img1, img2 = self.cropW() img11, img12 = img1.cropH() img21, img22 = img2.cropH() return img11, img12, img21, img22 def fpups(self): return AnnImg(self.img, self.ann.fpups()) def fconf(self, p): return AnnImg(self.img, self.ann.fconf(p)) def kconf(self, p): return AnnImg(self.img, self.ann.kconf(p)) def allNMS(self, th): return AnnImg(self.img, self.ann.allNMS(th)) @property def counts(self): return self.ann.counts @property def count(self): return self.ann.count def oneclass(self): return AnnImg(self.img, self.ann.oneclass()) def setbox(self, size): return AnnImg(self.img, self.ann.setbox(size)) @property def cxy(self): return self.ann[:, 0], self.ann[:, 1], self.ann[:, 2] def hm_pn(self, p=50, bg=None): aimg = self fp = aimg.hneg(th=True).fconf(p).ann #total = aimg.fcenter().ann total = aimg.fcenter().hmconf().ann tp = total.fconf(p) fn = total.kconf(p) fn.cl = 4 plot = None ann = tp.append(fn).append(fp) if bg is not None: plot = bg.wann(ann).plot(rect=False).img else: plot = aimg.wann(ann).plot(rect=False).img return ann, plot def vor(self, mind=None, maxd=None): return self.wann(self.ann.vor(mind, maxd)) def overlay(self, aimg): bg = self np1 = bg.np(t=False).copy() np2 = aimg.np().transpose((1, 2, 0)) c = 1 print(np1.shape, np2.shape) #np1 = np[:,:,c] np1[np2[:, :, 0] < 100] = (np1[np2[:, :, 0] < 100] * 0.5 + (np2[np2[:, :, 0] < 100]) * 0.5) return AnnImg(npimg=np1) @property def WH(self): return self.img.size @property def count(self): return self.ann.count def cropd(self, x, y, d): return self.crop(x, y, x + d, y + d) def cropdd(self, x, y, d): return self.crop(x - d, y - d, x + d, y + d) def rcropdd(self, d): W, H = self.WH return self.cropdd(randint(0, W), randint(0, H), d) def bcropW(aimg, d=100): W, H = aimg.WH xsort = np.sort(aimg.ann.xy[:, 0]) w = xsort[xsort.shape[0] // 2] return aimg.crop(0, 0, w + d, H), aimg.crop(w - d, 0, W, H), (w / W) def bcropH(aimg, d=100): W, H = aimg.WH ysort = np.sort(aimg.ann.xy[:, 1]) h = ysort[ysort.shape[0] // 2] return aimg.crop(0, 0, W, h + d), aimg.crop(0, h - d, W, H), (h / H) def bcrop(aimg, d=100): w1, w2, ww = aimg.bcropW(d) h1, h2, hh = aimg.bcropH(d) if (abs(0.5 - ww) < abs(0.5 - hh)): return w1, w2 else: return h1, h2 def collate_cls(self, cl, limit=8): ret = [] for c, ai in self.crop_ann(): if c == cl: ret.append(ai.np(t=False)) if len(ret) == 8: break if len(ret) == 0: return None return AnnImg(npimg=np.concatenate(ret, axis=1)).img def collate(self, limit=8): ret = [] for i in range(5): a = self.collate_cls(i, limit) if a is not None: ret.append((i, a)) return ret def nimg(self, d, ratio=1): numneg = int(self.count * ratio) W, H = self.WH arr = [] steps = 0 while len(arr) < numneg and steps < 10000: steps += 1 x, y = (randint(0, W), randint(0, H)) if self.cropdd(x, y, d).count == 0: arr.append((-1, x, y)) ann = Ann(dets=np.array(arr)) return AnnImg(self.img, ann) def crop_ann(self, d=None, cls=None): for cl, conf, xmin, ymin, xmax, ymax in self.ann.gen_dets(): if cls is not None and cls != cl: continue if d is not None: yield cl, self.cropdd((xmin + xmax) // 2, (ymin + ymax) // 2, d) else: yield cl, self.crop(xmin, ymin, xmax, ymax) def siamese(self, d, cl1, cl2, cnt=2): for clx, x in list(self.crop_ann(d, cl1)): self.ann.shuffle() for cly, y in islice(self.crop_ann(d, cl2), cnt): yield x, y, clx == cly def siamese23(self, d): return chain(self.siamese(d, 2, 2), self.siamese(d, 2, 3), self.siamese(d, 3, 2), self.siamese(d, 3, 3)) def blank(self, r=5, ratio=1): img = Image.new('RGBA', self.img.size, (0, 0, 0, 255)) if r is None: return AnnImg(img, self.ann).plot(rect=False, label=False, color=(255, 255, 255), bound=True, ratio=ratio) else: return AnnImg(img, self.ann).plotc(color=(255, 255, 255, 255), r=r) def clear(self, clr=(0, 0, 0, 255)): return AnnImg(Image.new('RGBA', self.img.size, (0, 0, 0, 255)), self.ann) def crop_neg(self, d, num): while num > 0: aimg = self.rcropdd(d) if aimg.count == 0: num -= 1 yield aimg def neg_mask(self, ratio=1): return AnnImg(npimg=self.blank(None, ratio=ratio).np()[0]) def mergeW(x, y): return AnnImg( npimg=np.concatenate([x.np(t=False), y.np(t=False)], axis=1)) def crop_posneg(self, d=None): cnt = 0 Wmax = 0 for cl, aimg in self.crop_ann(d): cnt += 1 W, H = aimg.WH Wmax = max(Wmax, W) yield cl, aimg for aimg in self.crop_neg(Wmax // 2, max(cnt, 50)): yield -1, aimg def cropc(self, d): x, y = self.WH x, y = x // 2, y // 2 return self.crop(x - d, y - d, x + d, y + d) def bound(self): xmin, ymin, xmax, ymax = self.ann.bound() return self.crop(xmin, ymin, xmax, ymax) def posmid(self, d, dist=None): if dist is None: dist = 1.5 * d pos = list([aimg.np() for cl, aimg in self.crop_ann(d)]) dets, _ = self.ann.midd(dist) ann = Ann(dets=dets) print("posmid", len(pos), ann.count) aimg = AnnImg(self.img, ann) mid = list([aimg.np() for cl, aimg in aimg.crop_ann(d)]) Xp, Xn = np.stack(pos), np.stack(mid) Yp, Yn = np.ones(Xp.shape[0]), np.zeros(Xn.shape[0]) return np.concatenate([Xp, Xn]), np.concatenate([Yp, Yn]), Xp, Xn def dbox(self, d): """ Increase the size of boxes by 2*d """ return self.wann(self.ann.dbox(d)) def posneg(self, d): pos = list([aimg.np() for cl, aimg in self.crop_ann(d)]) cl = list([cl for cl, aimg in self.crop_ann(d)]) #for i in range(len(pos)): neg = [] seed(0) while len(neg) < len(pos): aimg = self.rcropdd(d) if aimg.ann.count == 0: neg.append(aimg.np()) Xp, Xn = np.stack(pos), np.stack(neg) Yp, Yn = np.ones(Xp.shape[0]), np.zeros(Xn.shape[0]) Yp, Yn = np.array(cl) + 1, np.zeros(Xn.shape[0]) return np.concatenate([Xp, Xn]), np.concatenate([Yp, Yn]), Xp, Xn def posnegHW(self, d): i1, i2, i3, i4 = self.cropHW() X1, Y1, _, _ = i1.posneg(d) X2, Y2, _, _ = i2.posneg(d) X3, Y3, _, _ = i3.posneg(d) X4, Y4, _, _ = i4.posneg(d) Xt = np.concatenate([X1, X3]) Yt = np.concatenate([Y1, Y3]) Xv = np.concatenate([X2, X4]) Yv = np.concatenate([Y2, Y4]) return Xt, Yt, Xv, Yv def posmidHW(self, d): i1, i2, i3, i4 = self.cropHW() X1, Y1, _, _ = i1.posmid(d) X2, Y2, _, _ = i2.posmid(d) X3, Y3, _, _ = i3.posmid(d) X4, Y4, _, _ = i4.posmid(d) Xt = np.concatenate([X1, X3]) Yt = np.concatenate([Y1, Y3]) Xv = np.concatenate([X2, X4]) Yv = np.concatenate([Y2, Y4]) return Xt, Yt, Xv, Yv def tile(self, d, s): W, H = self.WH for x in range(0, W, s): for y in range(0, H, s): yield self.cropd(x, y, d), x, y def mtile(aimg, mimg, d, s, th=0.9): W, H = aimg.WH if mimg is None: mimg = AnnImg(npimg=np.zeros((W, H))) mimg = mimg.resize(W, H) for a, x, y in aimg.tile(d, s): m = mimg.cropd(x, y, d) p = 1 - (m.np().min() / 255.0) if (p > th): yield a, x, y #print("skip",x,y,p) def merge(self, fun, d=900, s=800, mimg=None): ret = [] for timg, x, y in self.mtile(mimg, d, s): ann = fun(timg) if ann is None: continue ann.dxy(-x, -y) ret.append(ann.dets) ret = np.concatenate(ret) return Ann(dets=ret) def test(self, model, d=900, s=800, cuda=True, mimg=None): model.eval() model.phase = "test" def fun(timg): img, tgt = timg.ssd() if cuda: img = img.cuda() dets = model.forward(Variable(img)) dim = img.size(2) dets = dets[0].data if cuda: dets = dets.cpu() ann = Ann(tensor=dets, dim=img.size(2)) return ann return self.merge(fun, d, s, mimg=mimg) def resize(self, x, y=None): if y is None: y = x W, H = self.img.size xs, ys = x / W, y / W ann = None if self.ann is not None: ann = self.ann.resize(x, y, W, H) return AnnImg(self.img.resize((x, y)), ann) def scale(self, s): x, y = self.WH return self.resize(int(x * s), int(y * s)) def pad(self, padx, pady=None): if pady is None: pady = padx W, H = self.WH return self.crop(-padx, -pady, W + padx, H + pady) def plt(self): img = self.img.copy() draw = ImageDraw.Draw(img) for c, x, y in zip(*self.cxy): draw.ellipse((x - d, y - d, x + d, y + d), fill=CLASS_COLORS[c]) return AnnImg(img, self.ann) def plot(self, ann=None, rect=True, label=True, color=None, save=None, r=5, bound=False, ratio=1): if self.ann is None: return AnnImg(self.img) if ann is not None: aimg = AnnImg( self.ann.plot(self.img, rect=False, text=label, color=color, r=r, bound=bound, ratio=ratio), ann) return aimg.plot() else: aimg = AnnImg( self.ann.plot(self.img, rect=rect, text=label, color=color, r=r, bound=bound, ratio=ratio), self.ann) if save is not None: aimg.img.save(save) return aimg def plotc(self, color=None, save=None, r=5, bound=False): if self.ann is None: return self.img aimg = self.plot(rect=False, label=False, color=color, r=r, bound=bound) return aimg.img def unet_mask(self, minr=0.5, maxr=1.5, vor=None): cimg = self img1 = cimg.neg_mask(ratio=minr) img2 = cimg.neg_mask(ratio=maxr) arr = (img1.np() + img2.np()) // 2 aimg = AnnImg(Image.fromarray(arr[0].astype('uint8'))).wann(self.ann) if vor is not None: aimg2 = AnnImg(aimg.plot_vor(width=vor)) if aimg2.np().sum() == 0: aimg2 = AnnImg(aimg.plot_vor(width=1)) aimg = aimg2 return aimg def _repr_png_(self): return self.plot().img._repr_png_() def save(self, img_path, ann_path=None): if ann_path is not None and self.ann is not None: self.ann.save(ann_path) self.img.save(img_path) def np(self, t=True): im = self.img arr = np.asarray(im, dtype=np.float) if len(arr.shape) == 3: if t: return np.transpose(arr, (2, 0, 1)) else: return arr else: return np.expand_dims(arr, 0) def negative(self, p=0.2): arr = self.np() arr = (arr < (p * 255)).astype('int32') img = Image.fromarray(np.uint8(arr[0] * 255)) return AnnImg(img, self.ann.copy()) def setScale(self, scale=50): return AnnImg(self.img, self.ann.setScale(scale)) def rescale(self, scale=50): ratio = scale / self.sc return self.scale(ratio) def ssd(aimg): tile, _ = aimg.WH target = None if aimg.ann is not None: target = aimg.ann.ssd(tile).float() img = aimg.np(t=False) height, width, _ = img.shape img -= (104, 117, 123) img = img.transpose(2, 0, 1) img = torch.from_numpy(img).squeeze().float() return img.unsqueeze(0), target def hist(self): img = self.np() ret = [] for i in range(img.shape[0]): hist, bins = np.histogram(img[i].flatten(), 256, [0, 256]) ret.append(hist) return np.concatenate(ret) def hist2(self): ret = [] W, H = self.WH for img in self.cropHW(): ret.append(img.hist()) return np.concatenate(ret) def dist(self, min=None, max=None): W, H = self.WH arr = self.ann.dist(W, H, min, max) img = Image.fromarray(arr.astype('uint8')) return AnnImg(img, self.ann) def neg(self, ratio=1.0, num=None, minn=None): if num is None: num = int(self.count * ratio) if minn is not None: num = max(num, minn) mask = self.neg_mask().inv().np() #print("mask", mask.shape) y, x = mask[0].nonzero() pts = np.stack([x, y], axis=1) np.random.shuffle(pts) xy = pts[:num] cl = -1 * np.ones((xy.shape[0], 1)) ann = Ann(dets=np.concatenate([cl, xy], axis=1).astype('int32')) return self.wann(ann) def hneg(self, ratio=1.0, num=None, minn=None, th=False, conf=10): if num is None: num = int(self.count * ratio) if minn is not None: num = max(num, minn) mask = self.neg_mask().inv().np() #print("mask", mask.shape) y, x = mask[0].nonzero() hm = self.np()[0] #print("hm", hm.shape) idx = np.argsort(hm[y, x]) py, px = y[idx], x[idx] xy = np.stack([px, py], axis=1) cl = -1 * np.ones((xy.shape[0], 1)) h = hm[xy[:, 1], xy[:, 0]] ann = Ann(dets=np.concatenate([cl.reshape(-1, 1), xy], axis=1).astype( 'int32')) ann.conf = ((255 - h) * 100) / 255 ann = ann.fconf(conf) #print(ann) if ann.count < num: return self.wann(ann) if th: ann = ann.pnms(50) return self.wann(ann.take(num)) def hmconf(self): return self.wann(self.ann.hmconf(self.np()[0])) def wann(self, ann): return AnnImg(self.img, ann) def append(self, ann): return AnnImg(self.img, self.ann.append(ann)) def inv(self): return AnnImg(npimg=(255 - self.np()[0]).astype('uint8'), ann=self.ann) def threshold(self, th): return AnnImg(npimg=(self.np()[0] < th) * 255) def watershed(self, th=128): distance1 = self.np()[0] W, H = self.WH markers1 = np.zeros((W, H)) for i, (x, y) in enumerate(self.ann.xy): markers1[y][x] = i image1 = distance1 < th labels1 = watershed(distance1, markers1, mask=image1) return labels1 def fcenter(self): W, H = self.WH return self.wann(self.ann.fcenter(0, 0, W, H)) def norm(self): npimg = self.np() mn = npimg.min() mx = npimg.max() out = 255 * (npimg - mn) / (mx - mn) return AnnImg(npimg=out[0], ann=self.ann) def cuts(self, labels1): XX = labels1.copy() X = labels1.copy() XX.fill(0) for i in range(1, X.shape[0] - 1): for j in range(1, X.shape[1] - 1): if X[i][j] != 0: if X[i][j] != X[i + 1][j] and X[i + 1][j] != 0: XX[i][j] = 1 if X[i][j] != X[i - 1][j] and X[i - 1][j] != 0: XX[i][j] = 1 if X[i][j] != X[i][j + 1] and X[i][j + 1] != 0: XX[i][j] = 1 if X[i][j] != X[i][j - 1] and X[i][j - 1] != 0: XX[i][j] = 1 return Image.fromarray(XX.astype('uint8') * 255) def plot_vor(cimg, color=0, width=4): #cimg = cimg.copy() points = cimg.ann.xy vor = Voronoi(points) center = vor.points.mean(axis=0) ptp_bound = vor.points.ptp(axis=0) finite_segments = [] for pointidx, simplex in zip(vor.ridge_points, vor.ridge_vertices): simplex = np.asarray(simplex) if np.all(simplex >= 0): finite_segments.append(vor.vertices[simplex]) else: i = simplex[simplex >= 0][0] # finite end Voronoi vertex t = vor.points[pointidx[1]] - vor.points[ pointidx[0]] # tangent t /= np.linalg.norm(t) n = np.array([-t[1], t[0]]) # normal midpoint = vor.points[pointidx].mean(axis=0) direction = np.sign(np.dot(midpoint - center, n)) * n far_point = vor.vertices[i] + direction * ptp_bound.max() finite_segments.append([vor.vertices[i], far_point]) #print(finite_segments) im = cimg.img.copy() draw = ImageDraw.Draw(im) for p1, p2 in finite_segments: draw.line((int(p1[0]), int(p1[1]), int(p2[0]), int(p2[1])), fill=color, width=width) return im