|
|
|
|
|
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) |
|
|
|
bb = BoundingBox(*cv2.boundingRect(coords)) |
|
rect = img[bb.y:bb.y+bb.h, bb.x:bb.x+bb.w] |
|
|
|
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) |
|
|
|
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)) |
|
|
|
|
|
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 |
|
|
|
|
|
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) |
|
|
|
|