import os import numpy as np import scipy.fft from scipy import ndimage import cv2 class Stitcher: def __init__(self,isstitch=False, n=[3,3], 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 = n[0] self.nrow = n[1] self.nominal_ov_pos=config_stitcher["nominal_ov_pos"] self.margin=config_stitcher["margin"] self.bound=self.nominal_ov_pos - self.margin # for sensor 1 self.pos_default = config_stitcher["pos_default"] 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) # 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) 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