spoof-detect / python /imtool.py
Niv Sardi
make watcher accept arguments,
26ef429
raw
history blame
7.4 kB
#!/usr/bin/env python3
import os
import math
import cv2
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
@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_anotation(self, id: int):
return f'{id} {self.x} {self.y} {self.w} {self.h}'
def read_marker(filename: str, Type: type):
ret = []
bco = None
with open(filename, 'r') as f:
lines = f.readlines()
for l in lines:
(b, x,y,w,h) = [float(i) for i in l.split(' ')]
bco = int(b)
ret.append(Type(x,y,w,h))
return bco, ret
def read_bounding_boxes(filename: str):
return read_marker(filename, BoundingBox)
def read_centroids(filename: str):
return read_marker(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)
x, y, w, h = cv2.boundingRect(coords) # Find minimum spanning bounding box
rect = img[y:y+h, x:x+w] # Crop the image - note we do this on the original image
return rect
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
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"{out}/debug"
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)