#!/usr/bin/env python3 import os import math import cv2 import base64 import numpy as np from typing import NamedTuple, Tuple, List from entity import Entity from common import mkdir TILE_SIZE = 416 TILE_OVERLAP = 0.8 class BoundingBox(NamedTuple): x: float = 0.0 y: float = 0.0 w: float = 0.0 h: float = 0.0 @classmethod def from_centroid(cls, c, shape): (ih, iw, ic) = shape self = cls(x=math.floor(w*(c.x - c.w/2)) , y=math.floor(h*(c.y - c.h/2)) , w=math.ceil(w*c.w) , h=math.ceil(h*c.h)) return self @classmethod def from_dict(cls, d): self = cls(x=d['x'], y=d['y'], w=d['width'], h=d['height']) return self @classmethod def from_arr(cls, a): self = cls(*a) return self @property def start(self): return floor_point(self.x, self.y) @property def end(self): return floor_point(self.x + self.w, self.y + self.h) def to_centroid(self, shape): (h, w, c) = shape return Centroid(x=math.floor(self.x + self.w/2)/w , y=math.floor(self.y + self.h/2)/h , w=math.ceil(self.w)/w , h=math.ceil(self.h)/h) def intersect(self, f, r: float = 0.8): six = self.x - f.x siy = self.y - f.y eix = six + self.w eiy = siy + self.h if six < 0: if six + self.w < 0: return None six = 0 if siy < 0: if siy + self.h < 0: return None siy = 0 if eix > f.w: if eix - self.w > f.w: return None eix = f.w if eiy > f.h: if eiy - self.h > f.h: return None eiy = f.h i = BoundingBox(six, siy, eix - six, eiy - siy) if (i.w*i.h) < (self.w*self.h)*r: return None return i class Centroid(BoundingBox): def to_bounding_box(self, shape): (h, w, c) = shape return BoundingBox( x=math.floor(w*(self.x - self.w/2)) , y=math.floor(h*(self.y - self.h/2)) , w=math.ceil(w*self.w) , h=math.ceil(h*self.h)) def to_annotation(self, id: int): return f'{id} {self.x} {self.y} {self.w} {self.h}' def read_base64(data): ib = base64.b64decode(data[22:]) arr = np.frombuffer(ib, dtype = np.uint8) return cv2.imdecode(arr, flags=cv2.IMREAD_COLOR) def read_markers(filename: str, Type: type): ret = [] with open(filename, 'r') as f: lines = f.readlines() for l in lines: try: (b, x,y,w,h, p) = [float(i) for i in l.split(' ')] except: try: (b, x,y,w,h) = [float(i) for i in l.split(' ')] except: continue p = -1 ret.append({"class": b, "prob": p, "box": Type(x,y,w,h)}) assert(len(ret)) return ret def read_centroids(filename: str): return read_markers(filename, Centroid) def coord_dict_to_point(c: dict): return coord_to_point(c['x'], c['y'], c['width'], c['height']) def coord_to_point(cx, cy, cw, ch): x = math.floor(cx + cw/2) y = math.floor(cy + ch/2) return f"{x} {y} {math.ceil(cw)} {math.ceil(ch)}" def floor_point(x: float, y: float): return (math.floor(x), math.floor(y)) def cut_img(im, s: Tuple[float, float], e: Tuple[float, float]): x1 = math.floor(s[0]) y1 = math.floor(s[1]) x2 = math.floor(e[0]) y2 = math.floor(e[1]) return im[y1:y2, x1:x2] def cut_logo(im, l): (x, y, w, h) = floor_logo(l) return im[y:y+h, x:x+w] def add_alpha(img): b, g, r = cv2.split(img) a = np.ones(b.shape, dtype=b.dtype) * 50 return cv2.merge((b,g,r,a)) def remove_white(img): gray = cv2.cvtColor(img, cv2.COLOR_BGRA2GRAY) gray = 255*(gray<128) coords = cv2.findNonZero(gray) # Find minimum spanning bounding box bb = BoundingBox(*cv2.boundingRect(coords)) rect = img[bb.y:bb.y+bb.h, bb.x:bb.x+bb.w] # Crop the image - note we do this on the original image return rect, bb def mix(a, b, fx, fy): alpha = b[:, :, 3]/255 return _mix_alpha(a, b, alpha, fx, fy) def mix_alpha(a, b, ba, fx, fy): (ah, aw, ac) = a.shape (bh, bw, bc) = b.shape p = 0.2 if (aw*p < bw or ah*p < bh): f = min(p*aw/bw, p*ah/bh) nw, nh = floor_point(bw*f, bh*f) #print(f'resizing to fit in {aw}x{ah}\t {bw}x{bh}\t=> {nw}x{nh}\tfactor {f}') r = cv2.resize(b, (nw, nh), interpolation = cv2.INTER_LINEAR) rba = cv2.resize(ba, (nw, nh), interpolation = cv2.INTER_LINEAR) return mix_alpha(a, r, rba, fx, fy) assert bw > 10, f'b({bw}) too small' assert bh > 10, f'b({bh}) too small' return _mix_alpha(a, b, ba, fx, fy) def _mix_alpha(a, b, ba, fx, fy): (ah, aw, ac) = a.shape (bh, bw, bc) = b.shape x = math.floor(fx*(aw - bw)) y = math.floor(fy*(ah - bh)) # handle transparency mat = a[y:y+bh,x:x+bw] cols = b[:, :, :3] mask = np.dstack((ba, ba, ba)) a[y:y+bh,x:x+bw] = mat * (1 - mask) + cols * mask #a[y:y+bh,x:x+bw] = cols return BoundingBox(x, y, bw, bh) def crop(id, fn, logos: List[Centroid], out = './data/squares'): basename = os.path.basename(fn).replace('.png', '') img_out = f"{out}/images" txt_out = f"{out}/labels" debug_out = f"{defaults.DEBUG_PATH}/{out}" mkdir.make_dirs([debug_out, img_out, txt_out]) im = cv2.imread(fn) rim = cv2.imread(fn) (h, w, c) = im.shape (tw, th) = (min(w, TILE_SIZE), min(h, TILE_SIZE)) (tx, ty)= ( math.ceil(w/(tw*TILE_OVERLAP)), math.ceil(h/(th*TILE_OVERLAP)) ) print('shape', basename, tx, ty, w, h) for x in range(tx): for y in range(ty): color = (0,x*(255/tx),y*(255/ty)) logo_color = (255, 0, 0) if tx < 2: xs = 0 else: xs = (w - tw)*x/(tx - 1) if ty < 2: ys = 0 else: ys = (h - th)*y/(ty - 1) f = BoundingBox(xs, ys, tw, th) start = floor_point(f.x, f.y) end = floor_point(f.x + f.w, f.y + f.h) rim = cv2.rectangle(rim, start, end, color, 10) li = [] for l in logos: bl = l.to_bounding_box(im.shape) rim = cv2.rectangle(rim, bl.start, bl.end, logo_color, 5) p = bl.intersect(f, 0.5) if p: li.append(p) nim = cut_img(im, start, end) rnim = cut_img(rim, start, end) img_name =f"{img_out}/{basename}-x{x}y{y}.jpg" txt_name =f"{txt_out}/{basename}-x{x}y{y}.txt" cv2.imwrite(img_name, nim) if len(li): with open(txt_name, 'w') as label: for p in li: dim = cv2.rectangle(rnim, p.start, p.end, logo_color, 5) lc = p.to_centroid((TILE_SIZE, TILE_SIZE, 3)) a = f"{int(id)} {lc.x} {lc.y} {lc.w} {lc.h}" label.write(a) cv2.imwrite(f'{debug_out}/{basename}{x}{y}.debug.png', dim) cv2.imwrite(f'{debug_out}/{basename}.debug.png', rim) if __name__ == '__main__': i = 0 with os.scandir('./data/') as it: for e in it: if e.name.endswith('.txt') and e.is_file(): print(e.name) try: i+=1 bco, boxes = read_bounding_boxes(e.path) crop(i, e.path.replace('.txt', '.png'), boxes) except Exception as err: print(err)