bluepulsar
588 lines
23 KiB
Python
588 lines
23 KiB
Python
import os
|
|
import numpy as np
|
|
import scipy.fft
|
|
from scipy import ndimage
|
|
import cv2
|
|
|
|
class Stitcher:
|
|
|
|
def __init__(self,isstitch=False, config_stitcher=None,verbose=False):
|
|
|
|
# Settings
|
|
self.isstitch = isstitch
|
|
self.config_stitcher = config_stitcher
|
|
self.verbose = verbose
|
|
# Find Boxes Related
|
|
self.min_area_frac = config_stitcher["min_area_frac"]
|
|
self.max_area_frac = config_stitcher["max_area_frac"]
|
|
self.max_aspect_ratio = config_stitcher["max_aspect_ratio"]
|
|
|
|
# stitch related
|
|
self.stitch_valid = False
|
|
self.boxes = self.units = None
|
|
self.index = config_stitcher["index"]
|
|
self.stitch_calibration = self.tile_list = []
|
|
self.ncol = config_stitcher["n"][0]
|
|
self.nrow = config_stitcher["n"][1]
|
|
self.nominal_ov_pos=config_stitcher["nominal_ov_pos"]
|
|
self.margin=config_stitcher["margin"]
|
|
self.bound=self.nominal_ov_pos - self.margin
|
|
|
|
if self.verbose:
|
|
print(f"Stitcher initialized with config: {self.config_stitcher}")
|
|
|
|
def stitch_rt_optimized(self, change_sv=False, s_mean=28., v_mean=180.):
|
|
tile_list = self.tile_list
|
|
stitch_calibration = self.stitch_calibration
|
|
ncol = self.ncol
|
|
nrow = self.nrow
|
|
|
|
processed_tiles = tile_list
|
|
img_stitched_height = None
|
|
|
|
for i in range(nrow):
|
|
# Start with first tile of the row
|
|
img_stitched_width = tile_list[i * ncol]
|
|
|
|
# Stitch horizontally within the row
|
|
for j in range(1, ncol):
|
|
img_1 = img_stitched_width # Already stitched part
|
|
img_2 = processed_tiles[i * ncol + j] # if change_sv else tile_list[i * ncol + j]
|
|
|
|
ov_pos = int(stitch_calibration[i * ncol + j - 1][1])
|
|
ov_pos, ov1, ov2 = self.__find_ov(img_1, img_2, ov_pos, dim=1)
|
|
|
|
# Use optimized stitch function
|
|
img_stitched_width = self.__stitch_optimized(img_1, img_2, ov_pos, ov1, ov2, dim=1, blend=True)
|
|
|
|
# Stitch rows vertically
|
|
if i == 0:
|
|
img_stitched_height = img_stitched_width
|
|
else:
|
|
# Handle width mismatch more efficiently
|
|
width_h = img_stitched_height.shape[1]
|
|
width_w = img_stitched_width.shape[1]
|
|
wdiff = width_h - width_w
|
|
|
|
if wdiff != 0:
|
|
# Use numpy pad instead of hstack for better performance
|
|
if wdiff > 0:
|
|
# Pad right side of the narrower image
|
|
padding = ((0, 0), (0, wdiff), (0, 0)) if len(img_stitched_width.shape) == 3 else ((0, 0), (0, wdiff))
|
|
img_stitched_width = np.pad(img_stitched_width, padding, mode='edge')
|
|
else:
|
|
# Crop the wider image
|
|
img_stitched_width = img_stitched_width[:, :width_h]
|
|
|
|
ov_pos_h = int(stitch_calibration[(i + 1) * nrow - 1][1])
|
|
ov_pos_h, ov_h1, ov_h2 = self.__find_ov(img_stitched_height, img_stitched_width, ov_pos_h, dim=0)
|
|
|
|
img_stitched_height = self.__stitch_optimized(
|
|
img_stitched_height, img_stitched_width, ov_pos_h, ov_h1, ov_h2, dim=0, blend=True
|
|
)
|
|
|
|
return img_stitched_height
|
|
|
|
def __find_ov(self, img1, img2, ov_pos, dim=1):
|
|
|
|
_dim = len(img1.shape)
|
|
|
|
if dim == 1:
|
|
if _dim == 3:
|
|
# RGB case
|
|
ov_1 = img1[:, -ov_pos:, :]
|
|
ov_2 = img2[:, :ov_pos, :]
|
|
else:
|
|
# gray scale
|
|
ov_1 = img1[:, -ov_pos:]
|
|
ov_2 = img2[:, :ov_pos]
|
|
elif dim == 0:
|
|
if _dim == 3:
|
|
ov_1 = img1[-ov_pos:, :, :]
|
|
ov_2 = img2[:ov_pos, :, :]
|
|
else:
|
|
# gray scale
|
|
ov_1 = img1[-ov_pos:, :]
|
|
ov_2 = img2[:ov_pos, :]
|
|
|
|
return ov_pos, ov_1, ov_2
|
|
|
|
def find_stitch_boxes(self,img_src,show = False):
|
|
|
|
img_corr = self.__correct_edge_bloom(img_src)
|
|
self.__detect_nonblack_regions(img_corr,show = show)
|
|
return len(self.boxes), self.boxes
|
|
|
|
# return self.boxes
|
|
|
|
|
|
def get_stitch_unit(self, img, flip_x=True, flip_y=True,index = np.arange(15)):
|
|
"""
|
|
Extract each tile given by boxes (x0,y0,x1,y1), default flip both X and Y due to inverting lens.
|
|
- img: source image (use corrected image if you prefer)
|
|
- boxes: list of (x0,y0,x1,y1)
|
|
- flip_x, flip_y: booleans (performed sequentially: X then Y)
|
|
- overlay_original: if True paste onto original image copy, otherwise onto blank background
|
|
- bg_color: background color for blank canvas (B,G,R)
|
|
Returns: canvas image and list of flipped tiles as (x0,y0,flipped_tile)
|
|
"""
|
|
h, w = img.shape[:2]
|
|
if self.units is None:
|
|
flipped_units = []
|
|
for i, (x0, y0, x1, y1) in enumerate(self.boxes):
|
|
|
|
# clamp coordinates
|
|
x0 = max(0, int(x0)); y0 = max(0, int(y0))
|
|
x1 = min(w, int(x1)); y1 = min(h, int(y1))
|
|
if x1 <= x0 or y1 <= y0:
|
|
continue
|
|
tile = img[y0:y1, x0:x1].copy()
|
|
# flip X then Y in default
|
|
if flip_x:
|
|
tile = cv2.flip(tile, 1) # horizontal
|
|
if flip_y:
|
|
tile = cv2.flip(tile, 0) # vertical
|
|
flipped_units.append((x0, y0, tile))
|
|
self.units = flipped_units
|
|
else:
|
|
flipped_units = self.units
|
|
index = self.index
|
|
for i, (x0, y0, x1, y1) in enumerate(self.boxes):
|
|
if i not in index:
|
|
continue
|
|
# clamp coordinates
|
|
x0 = max(0, int(x0)); y0 = max(0, int(y0))
|
|
x1 = min(w, int(x1)); y1 = min(h, int(y1))
|
|
if x1 <= x0 or y1 <= y0:
|
|
continue
|
|
tile = img[y0:y1, x0:x1].copy()
|
|
# flip X then Y in default
|
|
if flip_x:
|
|
tile = cv2.flip(tile, 1) # horizontal
|
|
if flip_y:
|
|
tile = cv2.flip(tile, 0) # vertical
|
|
# Assign instead of append
|
|
flipped_units[i] = (x0, y0, tile)
|
|
self.units = flipped_units
|
|
|
|
return True
|
|
|
|
def extract_tiles(self):
|
|
if not self.isstitch:
|
|
tile_idx = self.__get_tile_index_by_position(self.pos_default[7][0], self.pos_default[7][1])
|
|
if tile_idx is None:
|
|
tile_idx = 7
|
|
# Ensure index is always iterable (list or array)
|
|
self.index = [tile_idx] if isinstance(tile_idx, int) else tile_idx
|
|
self.ncol = 1
|
|
self.nrow = 1
|
|
|
|
tile_list = self.__get_tile_by_index(self.index)
|
|
self.tile_list = tile_list
|
|
|
|
return tile_list
|
|
|
|
def stitch_cal(self):
|
|
|
|
# init
|
|
tile_list = self.tile_list
|
|
img_stitched_height = None
|
|
stitch_calibration = []
|
|
# Flip col and row for porting
|
|
nrow, ncol = self.ncol, self.nrow
|
|
if not self.isstitch:
|
|
return self.tile_list[0], None
|
|
|
|
for i in range(ncol):
|
|
img_stitched_width = None
|
|
for j in range(nrow):
|
|
if j == 0:
|
|
img_j0 = tile_list[i*ncol]
|
|
img_stitched_width = img_j0
|
|
else:
|
|
img_1 = tile_list[i*ncol + j-1]
|
|
img_2 = tile_list[i*ncol + j]
|
|
|
|
ov_pos, ov1, ov2, c_corr = self.__get_ov(img_1.astype(np.float32), img_2.astype(np.float32), dim=1)
|
|
|
|
stitch_calibration.append(["hor {}".format(j-1),ov_pos])
|
|
|
|
# stitch images with stacking
|
|
img_stitched_width = self.__stitch_optimized(img_stitched_width, img_2, ov_pos, ov1, ov2, blend=True, dim=1)
|
|
|
|
if i == 0:
|
|
img_stitched_height = img_stitched_width
|
|
stitch_calibration.append(["ver {}".format(i),0])
|
|
pass
|
|
else:
|
|
width = img_stitched_height.shape[1]
|
|
width_new = img_stitched_width.shape[1]
|
|
wdiff = width - width_new
|
|
if wdiff > 0:
|
|
img_stitched_width = np.hstack((img_stitched_width, img_stitched_width[:, -wdiff:]))
|
|
elif wdiff < 0:
|
|
img_stitched_width = img_stitched_width[:, :wdiff]
|
|
|
|
ov_pos_h, ov_h1, ov_h2, c_corr = self.__get_ov(img_stitched_height, img_stitched_width, dim=0)
|
|
|
|
stitch_calibration.append(["ver {}".format(i),ov_pos_h])
|
|
|
|
img_stitched_height = self.__stitch_optimized(img_stitched_height, img_stitched_width, ov_pos_h, ov_h1, ov_h2, blend=True, dim=0)
|
|
|
|
self.stitch_calibration = stitch_calibration
|
|
|
|
return img_stitched_height, stitch_calibration
|
|
|
|
def __stitch_optimized(self, img1, img2, ov_pos, ov1, ov2, blend=True, dim=1):
|
|
_dim = len(img1.shape)
|
|
|
|
# Early return if no blending needed
|
|
if not blend:
|
|
ov = ov1
|
|
elif ov_pos <= 1:
|
|
ov = ov1
|
|
else:
|
|
# Vectorized blending
|
|
blend_coef = np.linspace(1.0, 0.0, ov_pos, dtype=np.float32)
|
|
|
|
# Prepare blend_coef for broadcasting
|
|
if dim == 1: # Horizontal blending
|
|
if _dim == 3:
|
|
# Shape: (1, ov_pos, 1)
|
|
blend_coef = blend_coef.reshape(1, ov_pos, 1)
|
|
else:
|
|
# Shape: (1, ov_pos)
|
|
blend_coef = blend_coef.reshape(1, ov_pos)
|
|
else: # Vertical blending
|
|
if _dim == 3:
|
|
# Shape: (ov_pos, 1, 1)
|
|
blend_coef = blend_coef.reshape(ov_pos, 1, 1)
|
|
else:
|
|
# Shape: (ov_pos, 1)
|
|
blend_coef = blend_coef.reshape(ov_pos, 1)
|
|
|
|
# Compute blended overlap in one operation
|
|
ov = (ov1.astype(np.float32) * blend_coef +
|
|
ov2.astype(np.float32) * (1 - blend_coef)).astype(np.uint8)
|
|
|
|
# Avoid unnecessary .copy() calls when possible
|
|
if dim == 1:
|
|
if _dim == 3:
|
|
# Use np.take or slicing without copy when safe
|
|
left = img1[:, :img1.shape[1] - ov_pos, :] if img1.shape[1] > ov_pos else img1[:, :0, :]
|
|
right = img2[:, ov_pos:, :]
|
|
stitched_img = np.concatenate((left, ov, right), axis=1)
|
|
else:
|
|
left = img1[:, :img1.shape[1] - ov_pos] if img1.shape[1] > ov_pos else img1[:, :0]
|
|
right = img2[:, ov_pos:]
|
|
stitched_img = np.concatenate((left, ov, right), axis=1)
|
|
else:
|
|
if _dim == 3:
|
|
top = img1[:img1.shape[0] - ov_pos, :, :] if img1.shape[0] > ov_pos else img1[:0, :, :]
|
|
bottom = img2[ov_pos:, :, :]
|
|
stitched_img = np.concatenate((top, ov, bottom), axis=0)
|
|
else:
|
|
top = img1[:img1.shape[0] - ov_pos, :] if img1.shape[0] > ov_pos else img1[:0, :]
|
|
bottom = img2[ov_pos:, :]
|
|
stitched_img = np.concatenate((top, ov, bottom), axis=0)
|
|
|
|
return stitched_img
|
|
|
|
def __get_ov(self, img1, img2, dim=1):
|
|
'''
|
|
default dimension is along axis 1 (width of the images)
|
|
'''
|
|
# pass parameters
|
|
nominal_ov = self.nominal_ov_pos
|
|
bound_ov = self.bound
|
|
|
|
cross_corr = []
|
|
h1, w1 = img1.shape[:2]
|
|
h2, w2 = img2.shape[:2]
|
|
# clamp tested overlap to valid range (1..min dimension)
|
|
max_possible = w1 if dim == 1 else h1
|
|
max_possible = min(max_possible, w2 if dim == 1 else h2)
|
|
lo = max(1, nominal_ov - bound_ov)
|
|
hi = min(max_possible, nominal_ov + bound_ov)
|
|
ov_space = range(lo, hi + 1)
|
|
|
|
for ov in ov_space:
|
|
if dim == 1:
|
|
ov_area1 = img1[:, max(0, w1 - ov):, :]
|
|
ov_area2 = img2[:, :min(ov, w2), :]
|
|
elif dim == 0:
|
|
ov_area1 = img1[max(0, h1 - ov):, :, :]
|
|
ov_area2 = img2[:min(ov, h2), :, :]
|
|
|
|
# require exact shape match (detect_nonblack_regions guarantees uniform tile size)
|
|
if ov_area1.size == 0 or ov_area2.size == 0 or ov_area1.shape != ov_area2.shape:
|
|
cross_corr.append(-1.0)
|
|
continue
|
|
|
|
# multi-channel cross correlation (average over channels)
|
|
corr = 0.0
|
|
for i in range(min(3, ov_area1.shape[2])):
|
|
a = ov_area1[:, :, i].astype(np.float32)
|
|
b = ov_area2[:, :, i].astype(np.float32)
|
|
denom = np.sqrt(np.sum(a * a) * np.sum(b * b)) + 1e-12
|
|
corr += float(np.sum(a * b) / denom)
|
|
corr = corr / float(min(3, ov_area1.shape[2]))
|
|
cross_corr.append(corr)
|
|
|
|
# pick best overlap (ignore invalid entries)
|
|
best_idx = int(np.argmax(cross_corr))
|
|
ov_pos = list(ov_space)[best_idx]
|
|
|
|
# extract the overlap regions (assume shapes match)
|
|
if dim == 1:
|
|
ov_1 = img1[:, max(0, w1 - ov_pos):, :].copy()
|
|
ov_2 = img2[:, :min(ov_pos, w2), :].copy()
|
|
else:
|
|
ov_1 = img1[max(0, h1 - ov_pos):, :, :].copy()
|
|
ov_2 = img2[:min(ov_pos, h2), :, :].copy()
|
|
return ov_pos, ov_1, ov_2, cross_corr
|
|
|
|
|
|
# Helper functions for stitching
|
|
def __get_tile_index_by_position(self, x, y, tol=20):
|
|
"""
|
|
Find flipped tile whose top-left matches (x,y) within tol pixels.
|
|
flipped_units: list of (x,y,tile)
|
|
"""
|
|
i= 0
|
|
for xi, yi, tile in self.units:
|
|
if abs(int(x) - int(xi)) <= tol and abs(int(y) - int(yi)) <= tol:
|
|
# return tile
|
|
return i
|
|
i += 1
|
|
return None
|
|
|
|
def __get_tile_by_index(self, idx):
|
|
"""
|
|
Retrieve tile by index
|
|
"""
|
|
if isinstance(idx, int):
|
|
idx = [idx]
|
|
|
|
if len(idx) > len(self.units):
|
|
raise ValueError("Index exceeds number of tiles.")
|
|
|
|
tile_list = []
|
|
for i in range(len(idx)):
|
|
[xi, yi, tile] = self.units[idx[i]]
|
|
tile_list.append(tile)
|
|
|
|
return tile_list
|
|
|
|
def __correct_edge_bloom(self, img, method="divide", kernel_size=None, clip=True):
|
|
"""
|
|
Simple illumination correction to reduce edge bloom / vignetting. Only for getting tile boxes.
|
|
- method: "divide" (default) or "subtract"
|
|
* divide: img_corrected = (img / (bg + eps)) * gain -- preserves colors
|
|
* subtract: img_corrected = img - (bg - mean_bg) -- removes low-freq background
|
|
- kernel_size: size for large-kernel blur used as background estimate (odd int).
|
|
If None, chosen from image size.
|
|
Returns corrected uint8 image.
|
|
"""
|
|
if img is None:
|
|
raise ValueError("Input image is None.")
|
|
|
|
h, w = img.shape[:2]
|
|
if kernel_size is None:
|
|
# choose large kernel relative to image size, ensure odd
|
|
k = max(31, min(h, w) // 2)
|
|
if k % 2 == 0:
|
|
k += 1
|
|
kernel_size = k
|
|
# convert to float32 for stable ops
|
|
img_f = img.astype(np.float32)
|
|
# estimate background by large Gaussian blur (works well for smooth bloom)
|
|
bg = cv2.GaussianBlur(img_f, (kernel_size, kernel_size), 0)
|
|
if method == "subtract":
|
|
# remove low-frequency background but keep overall brightness
|
|
mean_bg = bg.mean(axis=(0, 1), keepdims=True)
|
|
corr = img_f - (bg - mean_bg)
|
|
else:
|
|
# divide: normalized gain correction
|
|
eps = 1e-6
|
|
# compute gain so that mid-level stays near 128
|
|
mean_bg = bg.mean(axis=(0, 1), keepdims=True) + eps
|
|
gain = (mean_bg * 0.5) # target scale; 0.5 keeps values reasonable
|
|
corr = (img_f / (bg + eps)) * gain
|
|
if clip:
|
|
corr = np.clip(corr, 0, 255).astype(np.uint8)
|
|
else:
|
|
corr = corr.astype(np.uint8)
|
|
return corr
|
|
|
|
def __get_unit_pos(self,idx):
|
|
|
|
if isinstance(idx, int):
|
|
idx = [idx]
|
|
|
|
if len(idx) > len(self.units):
|
|
raise ValueError("Index exceeds number of vision units.")
|
|
|
|
pos_list = []
|
|
for i in range(len(idx)):
|
|
[xi, yi, tile] = self.units[idx[i]]
|
|
pos_list.append((xi, yi))
|
|
|
|
return pos_list
|
|
|
|
# Computation Functions
|
|
def __detect_nonblack_regions(self, img, morph_k=3, n_boxes_target=15, bk_thresh_step=10, bk_thresh_max=100,show = False):
|
|
"""
|
|
Return binary mask of non-black regions and list of bounding boxes (x,y,w,h).
|
|
Uses adaptive thresholding on grayscale + small morphology to be robust to illumination.
|
|
|
|
New behavior: after finding candidate boxes, center-crop all boxes to the same
|
|
width/height (smallest found) so returned boxes share a common size.
|
|
"""
|
|
if img is None:
|
|
raise ValueError("Input image is None.")
|
|
|
|
h, w = img.shape[:2]
|
|
n_boxes = 0
|
|
bk_thresh = 0 # initial black threshold
|
|
bk_thresh_step = 5
|
|
bk_thresh_max = 100
|
|
|
|
# Also include a low global V threshold to capture colored but dark tiles
|
|
hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)
|
|
v = hsv[:,:,2]
|
|
ww = hh = 0
|
|
while n_boxes != n_boxes_target and bk_thresh <= bk_thresh_max:
|
|
_, v_mask = cv2.threshold(v, bk_thresh, 255, cv2.THRESH_BINARY) # tune 50 if needed
|
|
|
|
mask = v_mask
|
|
|
|
# Morphology: remove small specks, fill small holes
|
|
k = max(1, morph_k)
|
|
kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (k, k))
|
|
mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel, iterations=1)
|
|
mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel, iterations=2)
|
|
# cv2.imshow("mask", mask)
|
|
# cv2.waitKey(0)
|
|
# find contours of candidate regions
|
|
contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
|
|
|
|
boxes = []
|
|
for c in contours:
|
|
x,y,ww,hh = cv2.boundingRect(c)
|
|
boxes.append((x,y,ww,hh))
|
|
|
|
# filter boxes by size (ignore too small or too large tiles)
|
|
boxes = self.__filter_boxes_by_size(boxes, (h, w))
|
|
n_boxes = len(boxes)
|
|
bk_thresh += bk_thresh_step
|
|
|
|
print("Detected {} candidate boxes with bk_thresh={}".format(n_boxes, bk_thresh - bk_thresh_step))
|
|
if n_boxes == 0:
|
|
self.boxes = []
|
|
return n_boxes
|
|
# compute common center-crop size: use minimum width/height among boxes
|
|
widths = [b[2] for b in boxes]
|
|
heights = [b[3] for b in boxes]
|
|
target_w = max(1,int(min(widths)))
|
|
target_h = max(1,int(min(heights)))
|
|
|
|
# center-crop each box to (target_w, target_h) inside original box bounds
|
|
cropped_boxes = []
|
|
for (x, y, ww, hh) in boxes:
|
|
cx = x + ww // 2
|
|
cy = y + hh // 2
|
|
nx0 = int(cx - target_w // 2)
|
|
ny0 = int(cy - target_h // 2)
|
|
# clamp to image bounds
|
|
nx0 = max(0, min(nx0, w - target_w))
|
|
ny0 = max(0, min(ny0, h - target_h))
|
|
# Convert to (x0, y0, x1, y1) format for direct image region retrieval
|
|
nx1 = nx0 + target_w
|
|
ny1 = ny0 + target_h
|
|
cropped_boxes.append((nx0, ny0, nx1, ny1))
|
|
|
|
# sort boxes left-to-right, top-to-bottom (useful for grid)
|
|
# Use quantized y-coordinates to group boxes into rows, making sorting robust to noise
|
|
if len(cropped_boxes) > 0:
|
|
# Calculate row tolerance based on typical box height (use median height)
|
|
heights = [b[3] - b[1] for b in cropped_boxes]
|
|
median_height = np.median(heights) if heights else target_h
|
|
# Use a fraction of the box height as tolerance (e.g., 20% of box height)
|
|
row_tolerance = max(1, int(median_height * 0.2))
|
|
|
|
# Quantize y-coordinates to group boxes into rows
|
|
def get_row_key(box):
|
|
y = box[1] # y0 coordinate
|
|
# Round y to nearest multiple of row_tolerance to group into rows
|
|
quantized_y = (y // row_tolerance) * row_tolerance
|
|
return (quantized_y, box[0]) # Sort by quantized row, then by x
|
|
|
|
cropped_boxes = sorted(cropped_boxes, key=get_row_key)
|
|
else:
|
|
cropped_boxes = sorted(cropped_boxes, key=lambda b: (b[1], b[0]))
|
|
|
|
self.boxes = cropped_boxes
|
|
img_with_boxes = self.draw_boxes(img= img)
|
|
if show:
|
|
cv2.imshow("Detected Boxes", img_with_boxes)
|
|
cv2.waitKey(0)
|
|
|
|
img_path = "./config/compound_eye.png"
|
|
if os.path.exists(img_path):
|
|
cv2.imwrite(img_path, img_with_boxes)
|
|
else:
|
|
os.makedirs(os.path.dirname(img_path), exist_ok=True)
|
|
cv2.imwrite(img_path, img_with_boxes)
|
|
# print("Check if compound eye detection image valid: {}".format(img_path))
|
|
return cropped_boxes
|
|
|
|
def draw_boxes(self, img, thickness=2):
|
|
"""
|
|
Draw bounding boxes on the image.
|
|
- img: input image (will be copied, not modified in-place)
|
|
- thickness: line thickness
|
|
Returns: image with drawn boxes
|
|
"""
|
|
if self.boxes is None or len(self.boxes) == 0:
|
|
print("No boxes to draw")
|
|
return img
|
|
|
|
img_draw = img.copy()
|
|
color=(0, 255, 0)
|
|
|
|
for i, (x0, y0, x1, y1) in enumerate(self.boxes):
|
|
# Draw rectangle
|
|
cv2.rectangle(img_draw, (x0, y0), (x1, y1), color, thickness)
|
|
|
|
# Draw index
|
|
text_size = cv2.getTextSize(str(i), cv2.FONT_HERSHEY_SIMPLEX, 0.6, 2)[0]
|
|
text_pos = (x0 + 5, y0 + text_size[1] + 5)
|
|
# Draw text background and text
|
|
cv2.rectangle(img_draw, (text_pos[0] - 2, text_pos[1] - text_size[1] - 2),
|
|
(text_pos[0] + text_size[0] + 2, text_pos[1] + 2), (0, 0, 0), -1)
|
|
cv2.putText(img_draw, str(i), text_pos, cv2.FONT_HERSHEY_SIMPLEX, 0.6, color, 2)
|
|
|
|
return img_draw
|
|
|
|
|
|
|
|
def __filter_boxes_by_size(self, boxes, img_shape):
|
|
"""
|
|
Remove boxes that are too small or too large relative to the image.
|
|
- boxes: list of (x,y,w,h)
|
|
- img_shape: (height, width)
|
|
- min_area_frac / max_area_frac: allowed area fraction of image
|
|
Returns filtered list of boxes.
|
|
"""
|
|
h, w = img_shape
|
|
img_area = h * w
|
|
min_area = img_area * float(self.min_area_frac)
|
|
max_area = img_area * float(self.max_area_frac)
|
|
filtered = []
|
|
for (x,y,ww,hh) in boxes:
|
|
area = ww * hh
|
|
if area < min_area or area > max_area or (abs(ww - hh) > self.max_aspect_ratio):
|
|
continue
|
|
filtered.append((x,y,ww,hh))
|
|
return filtered |