diff --git "a/app.py" "b/app.py" new file mode 100644--- /dev/null +++ "b/app.py" @@ -0,0 +1,2662 @@ +import os +import gc +import cv2 +import time +import tempfile +import mimetypes +import traceback +import numpy as np +import gradio as gr + +# --- Logging Helper --- +def log_and_print(message, current_log=""): + """Prints a message to the console and appends it to the log string.""" + print(message) # Print to console + return current_log + message + "\n" # Append to log string with newline + +# --- Helper Function: Crop Image by Percentage --- +def crop_image_by_percent(image, crop_top_percent=0.0, crop_bottom_percent=0.0): + """ + Crops the top and/or bottom portion of an image based on percentage. + + Args: + image: The input image (NumPy array). + crop_top_percent: Percentage of height to crop from the top (0-100). + crop_bottom_percent: Percentage of height to crop from the bottom (0-100). + + Returns: + The cropped image (NumPy array), or the original image if cropping is not needed + or percentages are invalid. Returns None if the input image is invalid. + """ + if image is None or image.size == 0: + # print("Warning: Invalid input image to crop_image_by_percent.") + return None # Return None for invalid input + + if crop_top_percent < 0 or crop_top_percent > 100 or \ + crop_bottom_percent < 0 or crop_bottom_percent > 100: + print(f"Warning: Invalid crop percentages ({crop_top_percent}%, {crop_bottom_percent}%). Must be between 0 and 100. Skipping crop.") + return image + + if crop_top_percent == 0 and crop_bottom_percent == 0: + return image # No cropping needed + + if crop_top_percent + crop_bottom_percent >= 100: + print(f"Warning: Total crop percentage ({crop_top_percent + crop_bottom_percent}%) is 100% or more. Skipping crop.") + return image + + try: + h, w = image.shape[:2] + + pixels_to_crop_top = int(h * crop_top_percent / 100.0) + pixels_to_crop_bottom = int(h * crop_bottom_percent / 100.0) + + start_row = pixels_to_crop_top + end_row = h - pixels_to_crop_bottom + + # Ensure indices are valid after calculation + if start_row >= end_row or start_row < 0 or end_row > h: + print(f"Warning: Invalid calculated crop rows (start={start_row}, end={end_row} for height={h}). Skipping crop.") + return image + + cropped_image = image[start_row:end_row, :] + # print(f"Debug: Cropped by percentage from {image.shape} to {cropped_image.shape}") + return cropped_image + except Exception as e: + print(f"Unexpected error during percentage cropping: {e}. Returning original image.") + traceback.print_exc() + return image + +# --- Helper Function: Crop Black Borders --- +def crop_black_borders(image, enable_cropping=True, strict_no_black_edges=False): + """ + Crops black borders from an image. + + Args: + image: The input image (NumPy array). + enable_cropping: If False, returns the original image. + strict_no_black_edges: If True, iteratively removes any remaining single black + pixel lines from the edges after the initial crop. + + Returns: + The cropped image (NumPy array), or the original image if cropping is disabled. + Returns None if the input is invalid or strict cropping removes everything. + """ + if not enable_cropping: + return image + if image is None or image.size == 0: + return None + + try: + # Check image channels before converting color + gray = None + mask_coords_found = False + coords = None + + # Attempt grayscale conversion first (common case) + if len(image.shape) == 3 and image.shape[2] == 3: + try: + gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY) + coords = cv2.findNonZero(gray) + if coords is not None: + mask_coords_found = True + except cv2.error as e_gray: + # print(f"Note: cvtColor to GRAY failed ({e_gray}), trying mask method.") + gray = None # Reset gray if conversion failed + elif len(image.shape) == 2: + gray = image # Already grayscale + coords = cv2.findNonZero(gray) + if coords is not None: + mask_coords_found = True + + # Fallback or alternative: Use mask if grayscale failed or shape is unusual + if not mask_coords_found: + try: + # Create a mask where any channel/value is > 0 + mask = np.any(image > 0, axis=-1) if len(image.shape) == 3 else (image > 0) + coords = cv2.findNonZero(mask.astype(np.uint8)) + if coords is not None: + mask_coords_found = True + except Exception as e_crop_fallback: + # print(f"Could not create mask for cropping fallback: {e_crop_fallback}. Returning original.") + return image # Cannot proceed if mask fails too + + if not mask_coords_found or coords is None: + # print("Debug: No non-black pixels found via any method, returning original.") + return image # Return original if all black or coords failed + + x, y, w, h = cv2.boundingRect(coords) + if w <= 0 or h <= 0: + # print(f"Debug: Invalid bounding rect ({w}x{h}), returning original.") + return image + + # Initial crop based on bounding rectangle + cropped_image = image[y:y+h, x:x+w] + + # --- START: Strict Edge Cropping Logic --- + if strict_no_black_edges and cropped_image is not None and cropped_image.size > 0: + # Iteratively remove black edges until none remain or image is empty + initial_shape = cropped_image.shape + iterations = 0 + MAX_ITERATIONS = max(initial_shape) # Safety break + while iterations < MAX_ITERATIONS: + iterations += 1 + # Re-check size in loop + if cropped_image is None or cropped_image.size == 0: + # print("Debug: Strict cropping resulted in empty image.") + return None # Image got cropped away entirely + + # Convert current crop to grayscale for edge checks + if len(cropped_image.shape) == 3: + if cropped_image.shape[2] == 1: # Handle case if somehow it becomes grayscale with 3 dims + gray_cropped = cropped_image[:, :, 0] + else: + try: + gray_cropped = cv2.cvtColor(cropped_image, cv2.COLOR_BGR2GRAY) + except cv2.error: + # print("Warning: Failed to convert to gray during strict crop, stopping strict loop.") + break # Stop if conversion fails + elif len(cropped_image.shape) == 2: + gray_cropped = cropped_image # Already grayscale + else: + # print("Warning: Unexpected image dimensions during strict crop, stopping strict loop.") + break # Stop if shape is weird + + # Check current edges + h_cr, w_cr = gray_cropped.shape[:2] + if h_cr <= 1 or w_cr <= 1: break # Cannot crop further if only 1 pixel wide/high + + top_row = gray_cropped[0, :] + bottom_row = gray_cropped[-1, :] + left_col = gray_cropped[:, 0] + right_col = gray_cropped[:, -1] + + top_has_black = np.any(top_row == 0) + bottom_has_black = np.any(bottom_row == 0) + left_has_black = np.any(left_col == 0) + right_has_black = np.any(right_col == 0) + + # If no edges have black pixels, we are done + if not (top_has_black or bottom_has_black or left_has_black or right_has_black): + # print(f"Debug: Strict cropping finished after {iterations-1} adjustments.") + break # Exit the while loop + + # Adjust cropping based on which edge(s) have black pixels + y_start_new, y_end_new = 0, h_cr + x_start_new, x_end_new = 0, w_cr + + if top_has_black: y_start_new += 1 + if bottom_has_black: y_end_new -= 1 + if left_has_black: x_start_new += 1 + if right_has_black: x_end_new -= 1 + + # Check if new bounds are valid before slicing + if y_start_new < y_end_new and x_start_new < x_end_new: + cropped_image = cropped_image[y_start_new:y_end_new, x_start_new:x_end_new] + else: + # print("Debug: Strict cropping bounds became invalid, stopping.") + cropped_image = None # Signal that cropping failed + break # Exit loop + + if iterations >= MAX_ITERATIONS: + print("Warning: Strict cropping reached max iterations, potential issue.") + if cropped_image is not None and initial_shape != cropped_image.shape: + print(f"Info: Strict cropping adjusted size from {initial_shape} to {cropped_image.shape}") + # --- END: Strict Edge Cropping Logic --- + + return cropped_image # Return the potentially strictly cropped image + + except cv2.error as e: + print(f"OpenCV Error during black border cropping: {e}. Returning uncropped image.") + return image + except Exception as e: + print(f"Unexpected error during black border cropping: {e}. Returning uncropped image.") + traceback.print_exc() + return image + +# --- Helper Function: Multi-Band Blending (Conceptual - Needs careful implementation) --- +def multi_band_blending(img1, img2, mask, num_levels=5): + # img1, img2: The two images to blend (float32, full canvas size) + # mask: The blending mask (float32, 0 to 1 transition, full canvas size, representing weight for img1) + # num_levels: Number of pyramid levels + log_message = "" # Add local logging if needed + + # Ensure inputs are float32 (caller should ensure this, but double check) + if img1.dtype != np.float32: img1 = img1.astype(np.float32) + if img2.dtype != np.float32: img2 = img2.astype(np.float32) + if mask.dtype != np.float32: + log_message = log_and_print(f"Warning: Mask input to multi_band_blending was {mask.dtype}, converting to float32.\n", log_message) + if mask.max() > 1: # Assuming uint8 if max > 1 + mask = mask.astype(np.float32) / 255.0 + else: # Assuming already float but maybe not float32 + mask = mask.astype(np.float32) + + # Ensure mask has same number of channels as images + if len(mask.shape) == 2 and len(img1.shape) == 3: + mask = cv2.cvtColor(mask, cv2.COLOR_GRAY2BGR) + elif len(mask.shape) == 3 and mask.shape[2] == 1 and len(img1.shape) == 3: + mask = cv2.cvtColor(mask, cv2.COLOR_GRAY2BGR) + elif mask.shape != img1.shape and mask.shape[:2] == img1.shape[:2]: + # If mask has 3 channels but img has 1 (unlikely but possible) + # Or other channel mismatches not covered above. Simpler to just ensure it matches. + log_message = log_and_print(f"Warning: Mask shape {mask.shape} mismatch with image shape {img1.shape}. Attempting replication.\n", log_message) + mask = cv2.cvtColor(cv2.cvtColor(mask, cv2.COLOR_BGR2GRAY), cv2.COLOR_GRAY2BGR) # Force to 3 channel based on img1 + + + # 1. Build Gaussian pyramids for img1, img2 + gp1 = [img1] + gp2 = [img2] + # Temporary list to store pyrDown results to avoid modifying list during iteration + gp1_next = [] + gp2_next = [] + actual_levels = 0 + for i in range(num_levels): + prev_h, prev_w = gp1[-1].shape[:2] + if prev_h < 2 or prev_w < 2: + log_message = log_and_print(f"Warning: Stopping image pyramid build at level {i} due to small size ({prev_h}x{prev_w}).\n", log_message) + break # Stop building pyramids for images + + try: + down1 = cv2.pyrDown(gp1[-1]) + down2 = cv2.pyrDown(gp2[-1]) + gp1_next.append(down1) + gp2_next.append(down2) + actual_levels += 1 # Increment count of successfully built levels + except cv2.error as e_pyrdown: + log_message = log_and_print(f"Error during pyrDown at level {i+1}: {e_pyrdown}. Stopping pyramid build.\n", log_message) + break # Stop if pyrDown fails + + # Update the main lists after the loop + gp1.extend(gp1_next); del gp1_next + gp2.extend(gp2_next); del gp2_next + gc.collect() + + # Adjust num_levels to the actual number built + num_levels = actual_levels + + # If pyramid build failed completely or input was too small + if num_levels == 0: + log_message = log_and_print("Error: Cannot build any pyramid levels. Using simple weighted average.\n", log_message) + blended_img = img1 * mask + img2 * (1.0 - mask) + blended_img = np.clip(blended_img, 0, 255).astype(np.uint8) + # print(log_message) # Optional: print warnings + if 'gp1' in locals(): del gp1 + if 'gp2' in locals(): del gp2 + gc.collect() + return blended_img # Fallback + + # 2. Build Laplacian pyramids for img1, img2 + # Smallest Gaussian level acts as base of Laplacian pyramid + lp1 = [gp1[num_levels]] + lp2 = [gp2[num_levels]] + for i in range(num_levels, 0, -1): + # Target size is the size of the *next larger* Gaussian level + target_size = (gp1[i-1].shape[1], gp1[i-1].shape[0]) + # log_message = log_and_print(f"Using resize instead of pyrUp for Laplacian level {i}\n", log_message) # Optional log + ge1 = cv2.resize(gp1[i], target_size, interpolation=cv2.INTER_LINEAR) + ge2 = cv2.resize(gp2[i], target_size, interpolation=cv2.INTER_LINEAR) + + # Ensure dimensions match EXACTLY before subtraction + # Sometimes pyrUp result might be 1 pixel off from the actual gp[i-1] size + h_target, w_target = gp1[i-1].shape[:2] + h_ge, w_ge = ge1.shape[:2] + + # Crop or pad ge1/ge2 to match gp1[i-1]/gp2[i-1] dimensions + if ge1.shape[:2] != (h_target, w_target): + #print(f"Level {i} pyrUp/resize shape mismatch: ge1={ge1.shape}, target={gp1[i-1].shape}. Adjusting ge1.") + ge1_adj = np.zeros_like(gp1[i-1], dtype=ge1.dtype) + copy_h = min(h_target, h_ge) + copy_w = min(w_target, w_ge) + ge1_adj[:copy_h, :copy_w] = ge1[:copy_h, :copy_w] + ge1 = ge1_adj + del ge1_adj + + if ge2.shape[:2] != (h_target, w_target): + #print(f"Level {i} pyrUp/resize shape mismatch: ge2={ge2.shape}, target={gp2[i-1].shape}. Adjusting ge2.") + ge2_adj = np.zeros_like(gp2[i-1], dtype=ge2.dtype) + copy_h = min(h_target, ge2.shape[0]) # Use ge2.shape[0] here + copy_w = min(w_target, ge2.shape[1]) # Use ge2.shape[1] here + ge2_adj[:copy_h, :copy_w] = ge2[:copy_h, :copy_w] + ge2 = ge2_adj + del ge2_adj + + + # Calculate Laplacian: Higher resolution Gaussian - Expanded lower resolution Gaussian + laplacian1 = cv2.subtract(gp1[i-1], ge1) + laplacian2 = cv2.subtract(gp2[i-1], ge2) + lp1.append(laplacian1) + lp2.append(laplacian2) + del ge1, ge2, laplacian1, laplacian2 + gc.collect() + + # del gp1, gp2 + # gc.collect() + + # lp1/lp2 lists are now [SmallestGaussian, LapN, LapN-1, ..., Lap1] (N=num_levels) + lp1.reverse() # Reverse to [Lap1, ..., LapN, SmallestGaussian] + lp2.reverse() + + # 3. Build Gaussian pyramid for the mask + gm = [mask] + gm_next = [] + actual_mask_levels = 0 + for i in range(num_levels): # Build mask pyramid only up to the actual image levels + prev_h, prev_w = gm[-1].shape[:2] + if prev_h < 2 or prev_w < 2: + log_message = log_and_print(f"Warning: Stopping mask pyramid build at level {i}.\n", log_message) + # num_levels should already be adjusted, but ensure mask levels don't exceed + break + try: + down_mask = cv2.pyrDown(gm[-1]) + gm_next.append(down_mask) + actual_mask_levels += 1 + except cv2.error as e_pyrdown_mask: + log_message = log_and_print(f"Error during mask pyrDown at level {i+1}: {e_pyrdown_mask}. Stopping mask pyramid build.\n", log_message) + break + + gm.extend(gm_next); del gm_next + gc.collect() + + # Ensure mask pyramid has the same number of levels as laplacian (+ base) + if len(gm) != num_levels + 1: + log_message = log_and_print(f"Error: Mask pyramid levels ({len(gm)}) does not match expected ({num_levels + 1}). Using simple average.\n", log_message) + # Fallback if mask pyramid construction failed unexpectedly + blended_img = img1 * mask + img2 * (1.0 - mask) + blended_img = np.clip(blended_img, 0, 255).astype(np.uint8) + if 'lp1' in locals(): del lp1 + if 'lp2' in locals(): del lp2 + if 'gm' in locals(): del gm + gc.collect() + return blended_img + + + # 4. Blend Laplacian levels + ls = [] # Blended Laplacian pyramid + for i in range(num_levels): # Blend Lap1 to LapN + lap1 = lp1[i] + lap2 = lp2[i] + mask_level = gm[i] # Use corresponding mask level (gm[0] for lp1[0]=Lap1, etc.) + + # Ensure mask shape matches laplacian shape for this level + if mask_level.shape[:2] != lap1.shape[:2]: + # print(f"Level {i} mask/lap shape mismatch: mask={mask_level.shape}, lap={lap1.shape}. Resizing mask.") + mask_level = cv2.resize(mask_level, (lap1.shape[1], lap1.shape[0]), interpolation=cv2.INTER_LINEAR) + # Ensure channels match after resize + if len(mask_level.shape) == 2 and len(lap1.shape) == 3: + mask_level = cv2.cvtColor(mask_level, cv2.COLOR_GRAY2BGR) + elif len(mask_level.shape) == 3 and mask_level.shape[2] == 1 and len(lap1.shape) == 3: + mask_level = cv2.cvtColor(mask_level, cv2.COLOR_GRAY2BGR) + # Clip mask just in case resize interpolation goes slightly out of [0,1] + mask_level = np.clip(mask_level, 0.0, 1.0) + + + # Blend: L = L1*Gm + L2*(1-Gm) + blended_lap = lap1 * mask_level + lap2 * (1.0 - mask_level) + ls.append(blended_lap) + del lap1, lap2, mask_level, blended_lap + gc.collect() + + # Blend the smallest Gaussian level (base of the pyramid) + base1 = lp1[num_levels] # Smallest Gaussian stored at the end of reversed lp1 + base2 = lp2[num_levels] + mask_base = gm[num_levels] # Use the smallest mask (corresponding to the smallest Gaussian level) + if mask_base.shape[:2] != base1.shape[:2]: + # print(f"Base level mask/base shape mismatch: mask={mask_base.shape}, base={base1.shape}. Resizing mask.") + mask_base = cv2.resize(mask_base, (base1.shape[1], base1.shape[0]), interpolation=cv2.INTER_LINEAR) + if len(mask_base.shape) == 2 and len(base1.shape) == 3: mask_base = cv2.cvtColor(mask_base, cv2.COLOR_GRAY2BGR) + elif len(mask_base.shape) == 3 and mask_base.shape[2]==1 and len(base1.shape) == 3: mask_base = cv2.cvtColor(mask_base, cv2.COLOR_GRAY2BGR) + mask_base = np.clip(mask_base, 0.0, 1.0) + + # Blend the base Gaussian level: B = B1*Gm_N + B2*(1-Gm_N) + blended_base = base1 * mask_base + base2 * (1.0 - mask_base) + ls.append(blended_base) # ls is now [BlendedLap1, ..., BlendedLapN, BlendedBase] + # del lp1, lp2, gm, base1, base2, mask_base, blended_base + del base1, base2, mask_base, blended_base + gc.collect() + + # 5. Reconstruct the final image from the blended Laplacian pyramid + # Start with the smallest blended base + blended_img = ls[num_levels] + for i in range(num_levels - 1, -1, -1): # Iterate from N-1 down to 0 + # Target size is the size of the *current* blended Laplacian level (ls[i]) + target_size = (ls[i].shape[1], ls[i].shape[0]) + # log_message = log_and_print(f"Using resize instead of pyrUp for reconstruction level {i}\n", log_message) # Optional log + expanded_prev = cv2.resize(blended_img, target_size, interpolation=cv2.INTER_LINEAR) + + # Delete previous level's blended_img (important for memory) + del blended_img + gc.collect() + + # Ensure dimensions match EXACTLY before adding + h_target_rec, w_target_rec = ls[i].shape[:2] + h_exp, w_exp = expanded_prev.shape[:2] + if expanded_prev.shape[:2] != (h_target_rec, w_target_rec): + # print(f"Reconstruction level {i} shape mismatch: expanded={expanded_prev.shape}, target={ls[i].shape}. Adjusting expanded.") + expanded_adj = np.zeros_like(ls[i], dtype=expanded_prev.dtype) + copy_h_rec = min(h_target_rec, h_exp) + copy_w_rec = min(w_target_rec, w_exp) + expanded_adj[:copy_h_rec, :copy_w_rec] = expanded_prev[:copy_h_rec, :copy_w_rec] + expanded_prev = expanded_adj + del expanded_adj + + # Add the blended Laplacian for the current level + current_laplacian = ls[i] # Get reference before add + blended_img = cv2.add(expanded_prev, current_laplacian) + del expanded_prev, current_laplacian # Remove laplacian reference ls[i] + ls[i] = None # Explicitly break the reference in the list too? Might help GC. + gc.collect() + + # Clip final result and convert back to uint8 + blended_img = np.clip(blended_img, 0, 255) + blended_img = blended_img.astype(np.uint8) + + # Optional: print warnings collected during the process + # if log_message: print("MultiBand Blend Logs:\n" + log_message) + + # Cleanup intermediate pyramids (important for memory) + del gp1, gp2, lp1, lp2, gm, ls + if 'laplacian1' in locals(): del laplacian1 + if 'laplacian2' in locals(): del laplacian2 + if 'ge1' in locals(): del ge1 + if 'ge2' in locals(): del ge2 + if 'mask_level' in locals(): del mask_level + if 'base1' in locals(): del base1 + if 'base2' in locals(): del base2 + if 'mask_base' in locals(): del mask_base + if 'blended_lap' in locals(): del blended_lap + if 'blended_base' in locals(): del blended_base + if 'expanded_prev' in locals(): del expanded_prev + gc.collect() + + return blended_img + +# --- Stitching Function: Focus on the pairwise images --- +def stitch_pairwise_images(img_composite, img_new, + transform_model_str="Homography", + blend_method="multi-band", + enable_gain_compensation=True, + orb_nfeatures=2000, + match_ratio_thresh=0.75, + ransac_reproj_thresh=5.0, + max_distance_coeff=0.5, + max_blending_width=10000, + max_blending_height=10000, + blend_smooth_ksize=15, + num_blend_levels=4 + ): + """ + Stitches a new image (img_new) onto an existing composite image (img_composite) + using an explicit, step-by-step pipeline (e.g., ORB features). + Allows choosing the geometric transformation model. + Returns the new composite. + """ + log_message = log_and_print("--- Starting pairwise stitch between composite and new image ---\n", "") + start_time_pairwise = time.time() + + # --- Input Validation --- + if img_composite is None or img_new is None: + log_message = log_and_print("Error: One or both input images are None for the pairwise stitching step.\n", log_message) + return None, log_message + if img_composite.size == 0 or img_new.size == 0: + log_message = log_and_print("Error: One or both input images are empty for the pairwise stitching step.\n", log_message) + return None, log_message + + h1, w1 = img_composite.shape[:2] + h2, w2 = img_new.shape[:2] + log_message = log_and_print(f"Pairwise Stitch: Img1({w1}x{h1}), Img2({w2}x{h2})\n", log_message) + log_message = log_and_print(f"Params: Transform={transform_model_str}, ORB Feats={orb_nfeatures}, Ratio Thresh={match_ratio_thresh}\n", log_message) + log_message = log_and_print(f"Params Cont'd: RANSAC Thresh={ransac_reproj_thresh}, Max Distance Coeff={max_distance_coeff}\n", log_message) + log_message = log_and_print(f"Blending: Method={blend_method}, GainComp={enable_gain_compensation}, SmoothKSize={blend_smooth_ksize}, MB Levels={num_blend_levels}\n", log_message) + + final_output_img = None # Initialize result variable + # Initialize other variables to None for better cleanup management + img1_u8, img2_u8 = None, None + kp1, des1, kp2, des2 = None, None, None, None + all_matches, good_matches = None, None + src_pts, dst_pts = None, None + H_matrix_3x3_for_canvas = None # Will hold the 3x3 matrix for canvas calculation (Affine or Homography) + final_warp_M = None # Will hold the actual 2x3 or 3x3 matrix for warping + mask_trans = None # Mask from estimation function (homography or affine) + pts1, dst_pts1_transformed = None, None + pts2, all_pts = None, None + output_img = None + warped_img1_u8 = None + mask_warped, mask_img2, overlap_mask = None, None, None + gain_applied_warped_img1_u8 = None + output_img_before_mb_float, blend_mask_float = None, None + img1_for_blend, img2_for_blend = None, None + is_affine = False # Flag to determine warp function + + try: + # --- Feature Detection and Matching --- + img1_u8 = img_composite.clip(0, 255).astype(np.uint8) if img_composite.dtype != np.uint8 else img_composite + img2_u8 = img_new.clip(0, 255).astype(np.uint8) if img_new.dtype != np.uint8 else img_new + + orb = cv2.ORB_create(nfeatures=orb_nfeatures) + kp1, des1 = orb.detectAndCompute(img1_u8, None) # keypoints and descriptors + kp2, des2 = orb.detectAndCompute(img2_u8, None) + + if des1 is None or des2 is None or len(kp1) < 2 or len(kp2) < 2: + log_message = log_and_print("Error: Not enough keypoints or descriptors found.\n", log_message) + if 'kp1' in locals(): del kp1 + if 'des1' in locals(): del des1 + if 'kp2' in locals(): del kp2 + if 'des2' in locals(): del des2 + del img1_u8, img2_u8 + gc.collect() + return None, log_message + log_message = log_and_print(f"Found {len(kp1)} keypoints in Img1, {len(kp2)} in Img2.\n", log_message) + + bf = cv2.BFMatcher(cv2.NORM_HAMMING, crossCheck=False) + # Check if descriptors are suitable for knnMatch (should be if ORB) + if des1.dtype != np.uint8: des1 = des1.astype(np.uint8) + if des2.dtype != np.uint8: des2 = des2.astype(np.uint8) + all_matches = bf.knnMatch(des1, des2, k=2) + del des1, des2; des1, des2 = None, None # Explicit delete + gc.collect() + + good_matches = [] + if all_matches is not None: + MAX_DISTANCE = max_distance_coeff * np.sqrt(w1**2 + h1**2) + # Filter out potential empty match pairs + valid_matches = [pair for pair in all_matches if isinstance(pair, (list, tuple)) and len(pair) == 2] + for m, n in valid_matches: + if m.distance < match_ratio_thresh * n.distance: + src_pt = np.array(kp1[m.queryIdx].pt) + dst_pt = np.array(kp2[m.trainIdx].pt) + distance = np.linalg.norm(dst_pt - src_pt) + + if distance < MAX_DISTANCE: + good_matches.append(m) + del valid_matches + del all_matches; all_matches = None + gc.collect() + + log_message = log_and_print(f"Found {len(good_matches)} good matches after ratio test.\n", log_message) + MIN_MATCH_COUNT = 10 # Keep a minimum threshold + + # --- Transformation Estimation (Homography or Affine) --- + if len(good_matches) >= MIN_MATCH_COUNT: + src_pts = np.float32([ kp1[m.queryIdx].pt for m in good_matches ]).reshape(-1,1,2) + dst_pts = np.float32([ kp2[m.trainIdx].pt for m in good_matches ]).reshape(-1,1,2) + del kp1, kp2, good_matches; kp1, kp2, good_matches = None, None, None # Explicit delete + gc.collect() + + estimation_failed = False + # Try Affine if selected + if transform_model_str == "Affine_Partial" or transform_model_str == "Affine_Full": + is_affine = True # Assume success initially + affine_matrix_2x3 = None + mask_a = None + try: + if transform_model_str == "Affine_Partial": + log_message = log_and_print(f"Attempting Affine Partial Estimation (RANSAC Thresh={ransac_reproj_thresh})...\n", log_message) + affine_matrix_2x3, mask_a = cv2.estimateAffinePartial2D(src_pts, dst_pts, method=cv2.RANSAC, ransacReprojThreshold=ransac_reproj_thresh) + else: # Affine_Full + log_message = log_and_print(f"Attempting Affine Full Estimation (RANSAC Thresh={ransac_reproj_thresh})...\n", log_message) + affine_matrix_2x3, mask_a = cv2.estimateAffine2D(src_pts, dst_pts, method=cv2.RANSAC, ransacReprojThreshold=ransac_reproj_thresh) + + if affine_matrix_2x3 is None: + raise ValueError(f"{transform_model_str} estimation returned None") + + # Convert 2x3 affine to 3x3 for canvas calculation consistency + H_matrix_3x3_for_canvas = np.vstack([affine_matrix_2x3, [0, 0, 1]]).astype(np.float64) + final_warp_M = affine_matrix_2x3.astype(np.float64) # Keep 2x3 for warpAffine + mask_trans = mask_a # Store the mask + + except Exception as e_affine: + log_message = log_and_print(f"Error during {transform_model_str} estimation: {e_affine}. Falling back to Homography.\n", log_message) + is_affine = False # Reset flag, will proceed to Homography block below + estimation_failed = True # Mark that the chosen affine failed + # Clean up affine specific vars if they exist + if 'affine_matrix_2x3' in locals(): del affine_matrix_2x3 + if 'mask_a' in locals(): del mask_a + H_matrix_3x3_for_canvas = None + final_warp_M = None + mask_trans = None + # NOTE: We are choosing to fall back instead of returning None immediately. + # If you prefer to fail hard if the selected affine fails, uncomment the next line: + # return None, log_message + + # Try Homography if selected OR if Affine failed and we are falling back + if not is_affine or estimation_failed: # If Homography was chosen or Affine failed + if estimation_failed: # Log if we are falling back + log_message = log_and_print("Falling back to Homography estimation...\n", log_message) + else: # Log if Homography was the original choice + log_message = log_and_print("Attempting Homography Estimation...\n", log_message) + + is_affine = False # Ensure flag is False for Homography path + H_matrix_homog = None + mask_h = None + try: + log_message = log_and_print(f"Estimating Homography (RANSAC Thresh={ransac_reproj_thresh})...\n", log_message) + H_matrix_homog, mask_h = cv2.findHomography(src_pts, dst_pts, cv2.RANSAC, ransacReprojThreshold=ransac_reproj_thresh) + if H_matrix_homog is None: + raise ValueError("Homography estimation returned None") + + H_matrix_3x3_for_canvas = H_matrix_homog.astype(np.float64) # Use this for canvas calc + final_warp_M = H_matrix_homog.astype(np.float64) # Use 3x3 for warpPerspective + mask_trans = mask_h # Store the mask + + except Exception as e_homog: + log_message = log_and_print(f"Error during Homography estimation: {e_homog}\n", log_message) + # Clean up if Homography itself fails + if 'H_matrix_homog' in locals(): del H_matrix_homog + if 'mask_h' in locals(): del mask_h + del src_pts, dst_pts + gc.collect() + return None, log_message # Fail if Homography (chosen or fallback) fails + + # --- Log Inliers from the successful estimation --- + model_name = "Affine" if is_affine else "Homography" + if mask_trans is not None: + inlier_count = np.sum(mask_trans) + log_message = log_and_print(f"{model_name} estimated with {inlier_count} inliers.\n", log_message) + if inlier_count < MIN_MATCH_COUNT: + log_message = log_and_print(f"Warning: Inlier count ({inlier_count}) < MIN_MATCH_COUNT for {model_name}. Result might be poor.\n", log_message) + del mask_trans; mask_trans = None # Delete the mask now + gc.collect() + else: + log_message = log_and_print(f"Warning: {model_name} mask was None.\n", log_message) + + + # --- Cleanup source/destination points --- + del src_pts, dst_pts; src_pts, dst_pts = None, None + gc.collect() + + # --- Canvas Calculation and Warping --- + pts1 = np.float32([[0,0],[0,h1-1],[w1-1,h1-1],[w1-1,0]]).reshape(-1,1,2) + try: + # Use the 3x3 matrix (derived from affine or directly from homography) for perspectiveTransform + # Ensure it's float64 + if H_matrix_3x3_for_canvas.dtype != np.float64: H_matrix_3x3_for_canvas = H_matrix_3x3_for_canvas.astype(np.float64) + dst_pts1_transformed = cv2.perspectiveTransform(pts1, H_matrix_3x3_for_canvas) + if dst_pts1_transformed is None: raise ValueError("perspectiveTransform returned None") + except Exception as e_tf: + model_name_tf = "Affine-derived" if is_affine else "Homography" + log_message = log_and_print(f"Error during perspectiveTransform (using {model_name_tf} 3x3 matrix): {e_tf}\n", log_message) + # Clean up before returning + del pts1 + if 'H_matrix_3x3_for_canvas' in locals(): del H_matrix_3x3_for_canvas + if 'final_warp_M' in locals(): del final_warp_M # Was holding the warp matrix + gc.collect() + return None, log_message + del pts1; pts1 = None + + pts2 = np.float32([[0,0],[0,h2-1],[w2-1,h2-1],[w2-1,0]]).reshape(-1,1,2) + # Ensure dst_pts1_transformed is float32 for concatenation if needed + all_pts = np.concatenate((pts2, dst_pts1_transformed.astype(np.float32)), axis=0) + del pts2, dst_pts1_transformed; pts2, dst_pts1_transformed = None, None + + padding = 2 + x_min, y_min = np.int32(all_pts.min(axis=0).ravel() - padding) + x_max, y_max = np.int32(all_pts.max(axis=0).ravel() + padding) + del all_pts; all_pts = None + gc.collect() + + translation_dist = [-x_min, -y_min] + H_translation = np.array([[1, 0, translation_dist[0]], [0, 1, translation_dist[1]], [0,0,1]], dtype=np.float64) + + output_width = x_max - x_min + output_height = y_max - y_min + + if output_width <= 0 or output_height <= 0 or output_width > max_blending_width or output_height > max_blending_height: + log_message = log_and_print(f"Error: Invalid output dimensions ({output_width}x{output_height}). Max allowed ({max_blending_width}x{max_blending_height})\n", log_message) + # Clean up before returning + if 'H_matrix_3x3_for_canvas' in locals(): del H_matrix_3x3_for_canvas + if 'final_warp_M' in locals(): del final_warp_M + if 'H_translation' in locals(): del H_translation + gc.collect() + return None, log_message + log_message = log_and_print(f"Calculated canvas size: {output_width}x{output_height}\n", log_message) + + # --- Memory Check for Blending --- + canvas_pixels = output_width * output_height + # Define a threshold based on available memory, e.g., 250 million pixels + # 15000*15000 = 225M, 30000*15000 = 450M + pixel_threshold = 225_000_000 + effective_blend_method = blend_method + + if blend_method == "multi-band" and canvas_pixels > pixel_threshold: + log_message = log_and_print(f"Warning: Canvas size ({output_width}x{output_height}, {canvas_pixels/1e6:.1f}M pixels) exceeds threshold ({pixel_threshold/1e6:.1f}M pixels) for multi-band blending.\n", log_message) + log_message = log_and_print("Switching to 'Linear' blending for this step to conserve memory.\n", log_message) + effective_blend_method = "linear" + + # Create output canvas + output_img = np.zeros((output_height, output_width, 3), dtype=np.uint8) + + # --- Calculate final transformation matrix for warping --- + # This incorporates the translation onto the canvas + final_warp_matrix_translated = None + if is_affine: + # We need the 2x3 matrix: (H_translation @ H_affine_3x3)[:2,:] + final_warp_matrix_translated = (H_translation @ H_matrix_3x3_for_canvas)[:2, :] + else: + # We need the 3x3 matrix: H_translation @ H_homography_3x3 + final_warp_matrix_translated = H_translation @ H_matrix_3x3_for_canvas # H_matrix_3x3 holds the homography here + + # --- Warp img1 onto the canvas --- + try: + if is_affine: + log_message = log_and_print("Warping image 1 using warpAffine...\n", log_message) + warped_img1_u8 = cv2.warpAffine(img1_u8, final_warp_matrix_translated, (output_width, output_height), flags=cv2.INTER_LINEAR, borderMode=cv2.BORDER_CONSTANT, borderValue=(0,0,0)) + else: + log_message = log_and_print("Warping image 1 using warpPerspective...\n", log_message) + warped_img1_u8 = cv2.warpPerspective(img1_u8, final_warp_matrix_translated, (output_width, output_height), flags=cv2.INTER_LINEAR, borderMode=cv2.BORDER_CONSTANT, borderValue=(0,0,0)) + except cv2.error as e_warp: + warp_type = 'Affine' if is_affine else 'Perspective' + log_message = log_and_print(f"Error during warping ({warp_type}): {e_warp}\n", log_message) + # Clean up before returning + if 'H_matrix_3x3_for_canvas' in locals(): del H_matrix_3x3_for_canvas + # final_warp_M was the matrix before translation + if 'final_warp_matrix_translated' in locals(): del final_warp_matrix_translated + if 'H_translation' in locals(): del H_translation + if 'img1_u8' in locals(): del img1_u8 + if 'output_img' in locals(): del output_img + gc.collect() + return None, log_message + + # --- Clean up matrices and source image --- + del H_matrix_3x3_for_canvas, H_translation, final_warp_matrix_translated, img1_u8 + # Note: final_warp_M (the untranslated matrix) is no longer needed + if 'final_warp_M' in locals(): del final_warp_M + gc.collect() + + # Place img2 onto the canvas + y_start, x_start = translation_dist[1], translation_dist[0] + y_end, x_end = y_start + h2, x_start + w2 + + # Define slicing for img2 read and canvas write, handling out-of-bounds placement + img2_y_start, img2_x_start = 0, 0 + img2_y_end, img2_x_end = h2, w2 + canvas_y_start, canvas_x_start = y_start, x_start + canvas_y_end, canvas_x_end = y_end, x_end + # Clip coordinates + if canvas_y_start < 0: img2_y_start = -canvas_y_start; canvas_y_start = 0 + if canvas_x_start < 0: img2_x_start = -canvas_x_start; canvas_x_start = 0 + if canvas_y_end > output_height: img2_y_end = h2 - (canvas_y_end - output_height); canvas_y_end = output_height + if canvas_x_end > output_width: img2_x_end = w2 - (canvas_x_end - output_width); canvas_x_end = output_width + + # Check if the calculated slices are valid + slice_h_canvas = canvas_y_end - canvas_y_start + slice_w_canvas = canvas_x_end - canvas_x_start + slice_h_img2 = img2_y_end - img2_y_start + slice_w_img2 = img2_x_end - img2_x_start + + mask_img2 = np.zeros(output_img.shape[:2], dtype=np.uint8) # Mask for img2 placement + img2_part = None + if slice_h_canvas > 0 and slice_w_canvas > 0 and slice_h_canvas == slice_h_img2 and slice_w_canvas == slice_w_img2: + img2_part = img2_u8[img2_y_start:img2_y_end, img2_x_start:img2_x_end] + output_img[canvas_y_start:canvas_y_end, canvas_x_start:canvas_x_end] = img2_part + mask_img2[canvas_y_start:canvas_y_end, canvas_x_start:canvas_x_end] = 255 + # Don't delete img2_part yet if needed for blend restoration + else: + log_message = log_and_print("Warning: Could not place img2 correctly onto the canvas.\n", log_message) + del img2_u8; img2_u8 = None # Input img2 no longer needed + gc.collect() + + # --- Create Masks for Blending --- + # Create mask for the warped image 1 (non-black pixels) + gray_warped = None + if warped_img1_u8 is not None: + gray_warped = cv2.cvtColor(warped_img1_u8, cv2.COLOR_BGR2GRAY) + # Ensure mask_warped is uint8 0 or 255 + # Check if gray_warped itself might be multi-channel if warp failed oddly? Should not happen. + if gray_warped is not None: + if len(gray_warped.shape) == 3: gray_warped = gray_warped[:,:,0] # Take one channel if needed + mask_warped = ((gray_warped > 0).astype(np.uint8)) * 255 + del gray_warped; gray_warped = None + gc.collect() + else: + mask_warped = np.zeros(output_img.shape[:2], dtype=np.uint8) # Empty mask if warp failed + + + # Find overlapping region mask (uint8 0 or 255) + overlap_mask = cv2.bitwise_and(mask_warped, mask_img2) + has_overlap = np.sum(overlap_mask > 0) > 0 # Check if any pixel > 0 + log_message = log_and_print(f"Overlap detected: {has_overlap}\n", log_message) + + # --- Gain Compensation --- + gain = 1.0 + gain_applied_warped_img1_u8 = warped_img1_u8 # Initialize with original warped image + + if enable_gain_compensation and has_overlap and warped_img1_u8 is not None: # Need warped image for gain comp + log_message = log_and_print("Gain Compensation Enabled. Calculating gain...\n", log_message) + try: + # --- Gain Calculation --- + gray_warped_for_gain = cv2.cvtColor(warped_img1_u8, cv2.COLOR_BGR2GRAY) + img2_gray = np.zeros_like(gray_warped_for_gain) + if slice_h_canvas > 0 and slice_w_canvas > 0: + if 0 <= canvas_y_start < canvas_y_end <= output_height and \ + 0 <= canvas_x_start < canvas_x_end <= output_width: + # Ensure output_img part is valid before cvtColor + img_to_convert = output_img[canvas_y_start:canvas_y_end, canvas_x_start:canvas_x_end] + if img_to_convert.size > 0: + img2_part_gray = cv2.cvtColor(img_to_convert, cv2.COLOR_BGR2GRAY) + img2_gray[canvas_y_start:canvas_y_end, canvas_x_start:canvas_x_end] = img2_part_gray + del img2_part_gray + else: log_message = log_and_print("Warning: Empty slice for gain calculation img2_gray.\n", log_message) + else: log_message = log_and_print("Warning: Invalid slice indices for gain calculation img2_gray.\n", log_message) + + overlap_mask_gain = overlap_mask # Use the already computed overlap mask + # Ensure masks are single channel before bitwise_and + if len(overlap_mask_gain.shape) == 3: overlap_mask_gain = overlap_mask_gain[:,:,0] + if len(gray_warped_for_gain.shape) == 3: gray_warped_for_gain = gray_warped_for_gain[:,:,0] + if len(img2_gray.shape) == 3: img2_gray = img2_gray[:,:,0] + + gray_warped_roi = cv2.bitwise_and(gray_warped_for_gain, gray_warped_for_gain, mask=overlap_mask_gain) + img2_roi = cv2.bitwise_and(img2_gray, img2_gray, mask=overlap_mask_gain) + del gray_warped_for_gain, img2_gray + + overlap_pixel_count = np.sum(overlap_mask_gain > 0) + if overlap_pixel_count > 0: + # Ensure ROIs are valid before calculating sum + mean1 = np.sum(gray_warped_roi[overlap_mask_gain > 0]) / overlap_pixel_count if gray_warped_roi is not None else 0 + mean2 = np.sum(img2_roi[overlap_mask_gain > 0]) / overlap_pixel_count if img2_roi is not None else 0 + + if mean1 > 1e-5 and mean2 > 1e-5: + gain = mean2 / mean1 + log_message = log_and_print(f"Calculated Gain: {gain:.2f}\n", log_message) + gain = np.clip(gain, 0.5, 2.0) # Clamp gain + log_message = log_and_print(f"Clamped Gain: {gain:.2f}\n", log_message) + else: + gain = 1.0 + log_message = log_and_print("Gain compensation skipped (means close to zero or invalid ROI).\n", log_message) + else: + gain = 1.0 + log_message = log_and_print("Gain compensation skipped (no overlap pixels).\n", log_message) + del gray_warped_roi, img2_roi + gc.collect() + # --- End Gain Calculation --- + + # Apply gain ONLY if calculated and different from 1.0 + if abs(gain - 1.0) > 1e-5: # Check float difference + gain_applied_float = warped_img1_u8.astype(np.float32) * gain + # *** Create new array for gain applied result *** + temp_gain_applied = gain_applied_float.clip(0, 255).astype(np.uint8) + # If gain_applied_warped_img1_u8 wasn't the original, delete it before reassigning + if gain_applied_warped_img1_u8 is not warped_img1_u8: + del gain_applied_warped_img1_u8 + gain_applied_warped_img1_u8 = temp_gain_applied # Assign the new gain-applied image + del gain_applied_float, temp_gain_applied + gc.collect() + log_message = log_and_print(f"Gain applied to warped image.\n", log_message) + else: + log_message = log_and_print("Gain is ~1.0, no gain applied.\n", log_message) + + except Exception as e_gain_calc: + gain = 1.0 + log_message = log_and_print(f"Warning: Error during gain calculation ({e_gain_calc}). Setting gain=1.0.\n", log_message) + # Ensure gain_applied remains the original warped image on error + if gain_applied_warped_img1_u8 is not warped_img1_u8: + del gain_applied_warped_img1_u8 # Delete potentially modified one + gc.collect() + gain_applied_warped_img1_u8 = warped_img1_u8 # Reset to original + # Clean up potential partial variables + if 'gray_warped_for_gain' in locals(): del gray_warped_for_gain + if 'img2_gray' in locals(): del img2_gray + if 'gray_warped_roi' in locals(): del gray_warped_roi + if 'img2_roi' in locals(): del img2_roi + if 'gain_applied_float' in locals(): del gain_applied_float + gc.collect() + elif warped_img1_u8 is None: + log_message = log_and_print("Skipping Gain Compensation as warped image is None.\n", log_message) + + # Ensure gain_applied_warped_img1_u8 holds the image to be used for blending + # (either original warped or gain-compensated version) + + # --- Blending Choice --- + # Blend using the potentially gain-compensated image: gain_applied_warped_img1_u8 + if effective_blend_method == "multi-band" and has_overlap and gain_applied_warped_img1_u8 is not None: + log_message = log_and_print(f"Applying Multi-band blending (Levels={num_blend_levels})...\n", log_message) + try: + # --- Generate Blend Mask using Distance Transform --- + log_message = log_and_print("Generating multi-band mask using distance transform...\n", log_message) + + # Ensure masks are single channel uint8 for distanceTransform + mask_warped_gray_mb = cv2.cvtColor(mask_warped, cv2.COLOR_BGR2GRAY) if len(mask_warped.shape) == 3 else mask_warped.copy() + mask_img2_gray_mb = cv2.cvtColor(mask_img2, cv2.COLOR_BGR2GRAY) if len(mask_img2.shape) == 3 else mask_img2.copy() + overlap_mask_gray_mb = cv2.cvtColor(overlap_mask, cv2.COLOR_BGR2GRAY) if len(overlap_mask.shape) == 3 else overlap_mask.copy() + + if mask_warped_gray_mb.dtype != np.uint8: mask_warped_gray_mb = (mask_warped_gray_mb > 0).astype(np.uint8) * 255 + if mask_img2_gray_mb.dtype != np.uint8: mask_img2_gray_mb = (mask_img2_gray_mb > 0).astype(np.uint8) * 255 + if overlap_mask_gray_mb.dtype != np.uint8: overlap_mask_gray_mb = (overlap_mask_gray_mb > 0).astype(np.uint8) * 255 + + + # Calculate distance transforms + # Distance to the nearest zero pixel (i.e., distance from the background) + dist1 = cv2.distanceTransform(mask_warped_gray_mb, cv2.DIST_L2, 5) + dist2 = cv2.distanceTransform(mask_img2_gray_mb, cv2.DIST_L2, 5) + + # Create float32 weight mask + weight1_norm = np.zeros(output_img.shape[:2], dtype=np.float32) + + # Identify non-overlapping regions (ensure using single channel masks) + non_overlap_mask1 = cv2.bitwise_and(mask_warped_gray_mb, cv2.bitwise_not(overlap_mask_gray_mb)) + non_overlap_mask2 = cv2.bitwise_and(mask_img2_gray_mb, cv2.bitwise_not(overlap_mask_gray_mb)) + + # Assign weights: 1.0 where only img1 exists, 0.0 where only img2 exists + weight1_norm[non_overlap_mask1 > 0] = 1.0 + weight1_norm[non_overlap_mask2 > 0] = 0.0 # Implicitly 0 initially, but good to be explicit + + # Calculate weights in the overlap region based on relative distance + # Weight for img1 = dist1 / (dist1 + dist2) + overlap_indices = np.where(overlap_mask_gray_mb > 0) + num_overlap_pixels = len(overlap_indices[0]) + if num_overlap_pixels > 0: + d1_overlap = dist1[overlap_indices] + d2_overlap = dist2[overlap_indices] + total_dist = d1_overlap + d2_overlap + # Avoid division by zero where total_dist is very small (deep inside both masks) + # If total_dist is near zero, assign weight based on which original mask was stronger? + # Using dist1 / (total_dist + epsilon) is simpler and generally works. + weights_overlap = d1_overlap / (total_dist + 1e-7) # Epsilon for stability + weight1_norm[overlap_indices] = np.clip(weights_overlap, 0.0, 1.0) + log_message = log_and_print(f"Calculated distance transform weights for {num_overlap_pixels} overlap pixels.\n", log_message) + else: + log_message = log_and_print("Warning: No overlap pixels found for distance transform weight calculation.\n", log_message) + + # Create boolean masks for later restoration steps + mask_warped_binary = (mask_warped_gray_mb > 0) + mask_img2_binary = (mask_img2_gray_mb > 0) + overlap_mask_binary = (overlap_mask_gray_mb > 0) + + # Clean up intermediate arrays from distance transform step + del mask_warped_gray_mb, mask_img2_gray_mb, overlap_mask_gray_mb, dist1, dist2 + del non_overlap_mask1, non_overlap_mask2 + del overlap_indices + if 'd1_overlap' in locals(): del d1_overlap + if 'd2_overlap' in locals(): del d2_overlap + if 'total_dist' in locals(): del total_dist + if 'weights_overlap' in locals(): del weights_overlap + gc.collect() + + # --- Apply Smoothing based on blend_smooth_ksize --- + blend_mask_float = weight1_norm # Start with the precise distance-based mask + if blend_smooth_ksize > 0 and blend_smooth_ksize % 2 == 1: + log_message = log_and_print(f"Smoothing multi-band blend mask with GaussianBlur ksize=({blend_smooth_ksize},{blend_smooth_ksize})...\n", log_message) + try: + # Need the boolean masks calculated above + + # Strict non-overlap areas (boolean arrays) + strict_non_overlap_mask1 = np.logical_and(mask_warped_binary, np.logical_not(overlap_mask_binary)) + strict_non_overlap_mask2 = np.logical_and(mask_img2_binary, np.logical_not(overlap_mask_binary)) + + + # Blur the original distance-based mask + weight1_norm_blurred = cv2.GaussianBlur(weight1_norm, (blend_smooth_ksize, blend_smooth_ksize), 0) + + # Clip the blurred mask to [0, 1] + blend_mask_float_blurred = np.clip(weight1_norm_blurred, 0.0, 1.0) + + # Assign the potentially blurred values first + blend_mask_float = blend_mask_float_blurred + + # Force 1.0 where only img1 should be + blend_mask_float[strict_non_overlap_mask1] = 1.0 + # Force 0.0 where only img2 should be + blend_mask_float[strict_non_overlap_mask2] = 0.0 + + log_message = log_and_print("Multi-band mask smoothed and edges restored.\n", log_message) + + except cv2.error as e_blur: + log_message = log_and_print(f"Warning: GaussianBlur failed for multi-band mask ({e_blur}). Using original distance-based mask.\n", log_message) + blend_mask_float = weight1_norm # Fallback to non-blurred + except Exception as e_blur_other: + log_message = log_and_print(f"Warning: Error during multi-band mask blur/restore ({e_blur_other}). Using original distance-based mask.\n", log_message) + blend_mask_float = weight1_norm # Fallback + finally: + # Clean up intermediate variables created in this block + if 'strict_non_overlap_mask1' in locals(): del strict_non_overlap_mask1 + if 'strict_non_overlap_mask2' in locals(): del strict_non_overlap_mask2 + if 'weight1_norm_blurred' in locals(): del weight1_norm_blurred + if 'blend_mask_float_blurred' in locals(): del blend_mask_float_blurred + gc.collect() + else: + log_message = log_and_print("Skipping multi-band mask smoothing (ksize not positive odd integer).\n", log_message) + # blend_mask_float is already weight1_norm (the precise one) + # --- End Smoothing --- + + # --- Prepare for Blending --- + img1_for_blend = gain_applied_warped_img1_u8.astype(np.float32) + # Store the state of output_img BEFORE multi-band blending + output_img_before_mb_float = output_img.astype(np.float32) + img2_for_blend = output_img_before_mb_float # Use the float version + + # --- Call Multi-Band Blending --- + blended_result_uint8 = multi_band_blending( + img1_for_blend, + img2_for_blend, + blend_mask_float, # The prepared mask + num_levels=num_blend_levels + ) + + # --- Restore Non-Overlap Regions --- + log_message = log_and_print("Restoring non-overlap regions after multi-band blending...\n", log_message) + + # Re-identify strict non-overlap boolean masks (using the ones calculated earlier) + strict_non_overlap_mask1 = np.logical_and(mask_warped_binary, np.logical_not(overlap_mask_binary)) + strict_non_overlap_mask2 = np.logical_and(mask_img2_binary, np.logical_not(overlap_mask_binary)) + + # Convert blended result to float for modification + output_img_float = blended_result_uint8.astype(np.float32) + + # Copy original pixels back into the non-overlap regions + # For img1's non-overlap region, use the (potentially gain compensated) warped img1 + output_img_float[strict_non_overlap_mask1] = img1_for_blend[strict_non_overlap_mask1] + + # For img2's non-overlap region, use the pixels from *before* blending + output_img_float[strict_non_overlap_mask2] = output_img_before_mb_float[strict_non_overlap_mask2] + + # Convert back to uint8 for the final result for this step + output_img = np.clip(output_img_float, 0, 255).astype(np.uint8) + log_message = log_and_print("Non-overlap regions restored.\n", log_message) + + # Optional final cleanup of absolute exterior (Post-blending mask) + combined_mask_binary = np.logical_or(mask_warped_binary, mask_img2_binary) + output_img[~combined_mask_binary] = 0 # Apply the sharp combined mask + log_message = log_and_print("Applied final exterior mask.\n", log_message) + + # Cleanup + del img1_for_blend, img2_for_blend, output_img_before_mb_float, blend_mask_float + del blended_result_uint8, output_img_float + del mask_warped_binary, mask_img2_binary, overlap_mask_binary + del strict_non_overlap_mask1, strict_non_overlap_mask2 + if 'combined_mask_binary' in locals(): del combined_mask_binary + if 'weight1_norm' in locals(): del weight1_norm + gc.collect() + log_message = log_and_print(f"Multi-band blending with restoration successful.\n", log_message) + + except Exception as e_blend: + log_message = log_and_print(f"Error during multi-band blending/restoration: {e_blend}. Falling back to simple overlay.\n{traceback.format_exc()}\n", log_message) + # Fallback uses gain_applied warped img1 over the original output_img + # Ensure mask_warped is usable by copyTo (needs same channel count or single channel) + mask_for_copy = mask_warped + if len(mask_warped.shape) == 2 and len(output_img.shape) == 3: + mask_for_copy = cv2.cvtColor(mask_warped, cv2.COLOR_GRAY2BGR) + elif len(mask_warped.shape) == 3 and len(output_img.shape) == 3 and mask_warped.shape[2] != output_img.shape[2]: + mask_for_copy = cv2.cvtColor(cv2.cvtColor(mask_warped, cv2.COLOR_BGR2GRAY), cv2.COLOR_GRAY2BGR) # Force 3 channels + + output_img = cv2.copyTo(gain_applied_warped_img1_u8, mask_for_copy, output_img) + if 'mask_for_copy' in locals() and mask_for_copy is not mask_warped: del mask_for_copy + + # Ensure cleanup if error happened mid-process + if 'img1_for_blend' in locals(): del img1_for_blend + if 'img2_for_blend' in locals(): del img2_for_blend + if 'output_img_before_mb_float' in locals(): del output_img_before_mb_float + if 'blend_mask_float' in locals(): del blend_mask_float + if 'blended_result_uint8' in locals(): del blended_result_uint8 + if 'mask_warped_binary' in locals(): del mask_warped_binary # Clean up boolean masks too + if 'mask_img2_binary' in locals(): del mask_img2_binary + if 'overlap_mask_binary' in locals(): del overlap_mask_binary + gc.collect() + + + # --- Linear Blending --- + elif effective_blend_method == "linear" and has_overlap and gain_applied_warped_img1_u8 is not None: + log_message = log_and_print("Applying Linear blending...\n", log_message) + # Ensure overlap_mask is single channel for findContours + overlap_mask_lin = cv2.cvtColor(overlap_mask, cv2.COLOR_BGR2GRAY) if len(overlap_mask.shape) == 3 else overlap_mask + contours, _ = cv2.findContours(overlap_mask_lin, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) + if not contours: + log_message = log_and_print("Warning: No contours in overlap. Using simple overlay.\n", log_message) + mask_for_copy = mask_warped # Prepare mask for copyTo + if len(mask_warped.shape) == 2 and len(output_img.shape) == 3: mask_for_copy = cv2.cvtColor(mask_warped, cv2.COLOR_GRAY2BGR) + elif len(mask_warped.shape) == 3 and len(output_img.shape) == 3 and mask_warped.shape[2] != output_img.shape[2]: mask_for_copy = cv2.cvtColor(cv2.cvtColor(mask_warped, cv2.COLOR_BGR2GRAY), cv2.COLOR_GRAY2BGR) + output_img = cv2.copyTo(gain_applied_warped_img1_u8, mask_for_copy, output_img) + if mask_for_copy is not mask_warped: del mask_for_copy + else: + main_contour = max(contours, key=cv2.contourArea) + x_overlap, y_overlap, w_overlap, h_overlap = cv2.boundingRect(main_contour) + # Clip bounding box to canvas dimensions + x_overlap = max(0, x_overlap); y_overlap = max(0, y_overlap) + w_overlap = min(w_overlap, output_width - x_overlap); h_overlap = min(h_overlap, output_height - y_overlap) + + if w_overlap <= 0 or h_overlap <= 0: + log_message = log_and_print("Warning: Invalid overlap bounding box after clipping. Using simple overlay.\n", log_message) + mask_for_copy = mask_warped # Prepare mask for copyTo + if len(mask_warped.shape) == 2 and len(output_img.shape) == 3: mask_for_copy = cv2.cvtColor(mask_warped, cv2.COLOR_GRAY2BGR) + elif len(mask_warped.shape) == 3 and len(output_img.shape) == 3 and mask_warped.shape[2] != output_img.shape[2]: mask_for_copy = cv2.cvtColor(cv2.cvtColor(mask_warped, cv2.COLOR_BGR2GRAY), cv2.COLOR_GRAY2BGR) + output_img = cv2.copyTo(gain_applied_warped_img1_u8, mask_for_copy, output_img) + if mask_for_copy is not mask_warped: del mask_for_copy + else: + # Create weight maps (float32) + weight1 = np.zeros(output_img.shape[:2], dtype=np.float32) + weight2 = np.zeros(output_img.shape[:2], dtype=np.float32) + blend_axis = 0 if w_overlap >= h_overlap else 1 + overlap_region_mask = overlap_mask_lin[y_overlap : y_overlap + h_overlap, x_overlap : x_overlap + w_overlap] + + # Generate gradient for the overlap box + gradient = None + if blend_axis == 0: # Horizontal blend + gradient = np.tile(np.linspace(1.0, 0.0, w_overlap, dtype=np.float32), (h_overlap, 1)) + else: # Vertical blend + gradient = np.tile(np.linspace(1.0, 0.0, h_overlap, dtype=np.float32).reshape(-1, 1), (1, w_overlap)) + + weight1_region = gradient + weight2_region = 1.0 - gradient + + # Apply weights only where the overlap mask is valid within the bounding box + valid_overlap = overlap_region_mask > 0 + weight1[y_overlap : y_overlap + h_overlap, x_overlap : x_overlap + w_overlap][valid_overlap] = weight1_region[valid_overlap] + weight2[y_overlap : y_overlap + h_overlap, x_overlap : x_overlap + w_overlap][valid_overlap] = weight2_region[valid_overlap] + del weight1_region, weight2_region, gradient, valid_overlap, overlap_region_mask + gc.collect() + + # Assign weights for non-overlapping regions (ensure masks are single channel) + mask_warped_lin = cv2.cvtColor(mask_warped, cv2.COLOR_BGR2GRAY) if len(mask_warped.shape) == 3 else mask_warped + mask_img2_lin = cv2.cvtColor(mask_img2, cv2.COLOR_BGR2GRAY) if len(mask_img2.shape) == 3 else mask_img2 + non_overlap_mask1 = cv2.bitwise_and(mask_warped_lin, cv2.bitwise_not(overlap_mask_lin)) + weight1[non_overlap_mask1 > 0] = 1.0 + non_overlap_mask2 = cv2.bitwise_and(mask_img2_lin, cv2.bitwise_not(overlap_mask_lin)) + weight2[non_overlap_mask2 > 0] = 1.0 # Weight for image 2 is 1.0 in its non-overlap area + + # Normalize weights before potential smoothing + total_weight = weight1 + weight2 + 1e-6 # Add epsilon + weight1_norm = weight1 / total_weight + weight2_norm = weight2 / total_weight + del weight1, weight2, total_weight + gc.collect() + + # --- Apply Smoothing based on blend_smooth_ksize --- + if blend_smooth_ksize > 0 and blend_smooth_ksize % 2 == 1: + log_message = log_and_print(f"Smoothing linear blend weights with GaussianBlur ksize=({blend_smooth_ksize},{blend_smooth_ksize})...\n", log_message) + try: + # Identify the actual blending area (where both weights contribute meaningfully and overlap exists) + overlap_area_mask_bool = (weight1_norm > 1e-6) & (weight2_norm > 1e-6) & (overlap_mask_lin > 0) + + smoothed_w1 = cv2.GaussianBlur(weight1_norm, (blend_smooth_ksize, blend_smooth_ksize), 0) + smoothed_w2 = cv2.GaussianBlur(weight2_norm, (blend_smooth_ksize, blend_smooth_ksize), 0) + + # Renormalize smoothed weights ONLY in the overlap area + total_smoothed_weight = smoothed_w1 + smoothed_w2 + 1e-6 + # Use temporary arrays to avoid modifying originals during calculation if needed + temp_w1 = weight1_norm.copy() # Work on copies + temp_w2 = weight2_norm.copy() + temp_w1[overlap_area_mask_bool] = (smoothed_w1 / total_smoothed_weight)[overlap_area_mask_bool] + temp_w2[overlap_area_mask_bool] = (smoothed_w2 / total_smoothed_weight)[overlap_area_mask_bool] + + # Restore strict 1.0 / 0.0 weights in non-overlap areas + temp_w1[ non_overlap_mask1 > 0 ] = 1.0 + temp_w1[ non_overlap_mask2 > 0 ] = 0.0 + temp_w2[ non_overlap_mask1 > 0 ] = 0.0 + temp_w2[ non_overlap_mask2 > 0 ] = 1.0 + + # Assign back to the working variables + weight1_norm = temp_w1 + weight2_norm = temp_w2 + + del smoothed_w1, smoothed_w2, total_smoothed_weight, overlap_area_mask_bool, temp_w1, temp_w2 + gc.collect() + log_message = log_and_print("Linear weights smoothed and renormalized.\n", log_message) + + except cv2.error as e_blur: + log_message = log_and_print(f"Warning: GaussianBlur failed for linear weights ({e_blur}). Using original weights.\n", log_message) + except Exception as e_blur_other: + log_message = log_and_print(f"Warning: Error during linear weight smoothing ({e_blur_other}). Using original weights.\n", log_message) + finally: + # Ensure cleanup of temp vars in this block + if 'smoothed_w1' in locals(): del smoothed_w1 + if 'smoothed_w2' in locals(): del smoothed_w2 + if 'total_smoothed_weight' in locals(): del total_smoothed_weight + if 'overlap_area_mask_bool' in locals(): del overlap_area_mask_bool + if 'temp_w1' in locals(): del temp_w1 + if 'temp_w2' in locals(): del temp_w2 + gc.collect() + else: + log_message = log_and_print("Skipping linear weight smoothing (ksize not positive odd integer).\n", log_message) + # --- End Smoothing --- + + + # Blend using potentially smoothed and renormalized weights + # Identify regions: where img1 only, img2 only, and blend region + non_overlap_mask1_bool = (non_overlap_mask1 > 0) + non_overlap_mask2_bool = (non_overlap_mask2 > 0) + blend_mask_bool = np.logical_not(np.logical_or(non_overlap_mask1_bool, non_overlap_mask2_bool)) & (overlap_mask_lin > 0) + + # Copy non-overlapping part of image 1 directly where its weight is 1 + output_img[non_overlap_mask1_bool] = gain_applied_warped_img1_u8[non_overlap_mask1_bool] + + # Non-overlapping part of image 2 is already in output_img from the initial placement + + # Blend the overlapping/transition areas + blend_indices = np.where(blend_mask_bool) + num_blend_pixels = len(blend_indices[0]) + + if num_blend_pixels > 0: + log_message = log_and_print(f"Blending {num_blend_pixels} pixels linearly...\n", log_message) + try: + # Ensure images are float32 for blending calculation + img1_blend_float = gain_applied_warped_img1_u8[blend_indices].astype(np.float32) + img2_blend_float = output_img[blend_indices].astype(np.float32) # Pixels already placed from img2 + + # Get weights for the blend region and broadcast for element-wise multiplication + w1_blend_1d = weight1_norm[blend_indices] + w2_blend_1d = weight2_norm[blend_indices] + # Add new axis for broadcasting: (N,) -> (N, 1) to multiply with (N, 3) pixel data + w1_blend_broadcast = w1_blend_1d[:, np.newaxis] + w2_blend_broadcast = w2_blend_1d[:, np.newaxis] + + # Perform the weighted sum + blended_float = w1_blend_broadcast * img1_blend_float + w2_blend_broadcast * img2_blend_float + blended_uint8 = blended_float.clip(0, 255).astype(np.uint8) + + # Place the blended result back into the output image + output_img[blend_indices] = blended_uint8 + + del img1_blend_float, img2_blend_float, w1_blend_1d, w2_blend_1d + del w1_blend_broadcast, w2_blend_broadcast, blended_float, blended_uint8 + gc.collect() + log_message = log_and_print("Linear blending successful.\n", log_message) + + except MemoryError: + log_message = log_and_print("Warning: MemoryError during float blending. Using simple overlay for blend region.\n", log_message) + # Fallback: copy img1 over img2 in the blend region + blend_mask_uint8 = blend_mask_bool.astype(np.uint8) * 255 + mask_for_copy = cv2.cvtColor(blend_mask_uint8, cv2.COLOR_GRAY2BGR) if len(output_img.shape) == 3 else blend_mask_uint8 + if np.any(mask_for_copy): + output_img = cv2.copyTo(gain_applied_warped_img1_u8, mask_for_copy, output_img) + del blend_mask_uint8, mask_for_copy + gc.collect() + except Exception as e_blend_lin: + log_message = log_and_print(f"Warning: Error during float blending ({e_blend_lin}). Using simple overlay for blend region.\n", log_message) + blend_mask_uint8 = blend_mask_bool.astype(np.uint8) * 255 + mask_for_copy = cv2.cvtColor(blend_mask_uint8, cv2.COLOR_GRAY2BGR) if len(output_img.shape) == 3 else blend_mask_uint8 + if np.any(mask_for_copy): + output_img = cv2.copyTo(gain_applied_warped_img1_u8, mask_for_copy, output_img) + del blend_mask_uint8, mask_for_copy + gc.collect() + else: + log_message = log_and_print("Note: Linear blend mask was empty, skipping float blend step.\n", log_message) + + # Clean up linear blending specific variables + del weight1_norm, weight2_norm, blend_mask_bool + del non_overlap_mask1, non_overlap_mask2, non_overlap_mask1_bool, non_overlap_mask2_bool + del mask_warped_lin, mask_img2_lin, overlap_mask_lin + if 'blend_indices' in locals(): del blend_indices + gc.collect() + + # Clean up contour variables regardless of path taken inside linear blend + if 'contours' in locals(): del contours + if 'main_contour' in locals(): del main_contour + gc.collect() + + # Simple overlay if no blending applied or specified OR if warped image was None + elif not has_overlap or effective_blend_method not in ["linear", "multi-band"] or gain_applied_warped_img1_u8 is None: + if gain_applied_warped_img1_u8 is None: + log_message = log_and_print("Warped image was None. Performing simple overlay (only showing img2).\n", log_message) + # In this case, output_img already contains img2 where it should be, and black elsewhere. + # No copyTo needed, as there's nothing to copy from. + elif not has_overlap: + log_message = log_and_print("No overlap. Performing simple overlay.\n", log_message) + else: + log_message = log_and_print(f"Blending method '{effective_blend_method}' or overlap condition not met. Performing simple overlay.\n", log_message) + + if gain_applied_warped_img1_u8 is not None: # Only copy if we have something to copy + # Overlay gain_applied warped img1 onto output_img where mask_warped is non-zero + mask_for_copy = mask_warped # Prepare mask for copyTo + if len(mask_warped.shape) == 2 and len(output_img.shape) == 3: mask_for_copy = cv2.cvtColor(mask_warped, cv2.COLOR_GRAY2BGR) + elif len(mask_warped.shape) == 3 and len(output_img.shape) == 3 and mask_warped.shape[2] != output_img.shape[2]: mask_for_copy = cv2.cvtColor(cv2.cvtColor(mask_warped, cv2.COLOR_BGR2GRAY), cv2.COLOR_GRAY2BGR) + output_img = cv2.copyTo(gain_applied_warped_img1_u8, mask_for_copy, output_img) + if mask_for_copy is not mask_warped: del mask_for_copy + + # --- Final Result Assignment --- + final_output_img = output_img # Assign the final blended/overlaid image + + end_time_pairwise = time.time() + log_message = log_and_print(f"Pairwise stitching finished. Time: {end_time_pairwise - start_time_pairwise:.2f}s\n", log_message) + + else: # Not enough good matches + log_message = log_and_print(f"Error: Not enough good matches ({len(good_matches)} < {MIN_MATCH_COUNT}).\n", log_message) + # Minimal cleanup needed here, mostly handled in finally block + if 'kp1' in locals(): del kp1 + if 'kp2' in locals(): del kp2 + if 'good_matches' in locals(): del good_matches + + + except Exception as e: + log_message = log_and_print(f"Error during pairwise stitching: {e}\n{traceback.format_exc()}\n", log_message) + final_output_img = None # Ensure None is returned on error + + finally: + # --- Comprehensive Cleanup --- + # Delete variables in roughly reverse order of creation / dependency + # Blend-specific intermediates + if 'img1_for_blend' in locals(): del img1_for_blend + if 'img2_for_blend' in locals(): del img2_for_blend + if 'output_img_before_mb_float' in locals(): del output_img_before_mb_float + if 'blend_mask_float' in locals(): del blend_mask_float + if 'weight1_norm' in locals(): del weight1_norm # From mask gen (MB or Linear) + if 'weight2_norm' in locals(): del weight2_norm # From Linear mask gen + # ... other linear/MB intermediate vars ... + + # Gain/Warp intermediates + if 'gain_applied_warped_img1_u8' in locals() and gain_applied_warped_img1_u8 is not None: + # Only delete if it's a separate copy from warped_img1_u8 + if 'warped_img1_u8' in locals() and warped_img1_u8 is not None and gain_applied_warped_img1_u8 is not warped_img1_u8: + del gain_applied_warped_img1_u8 + # else it points to warped_img1_u8 or warped_img1_u8 is None/deleted already + + if 'warped_img1_u8' in locals() and warped_img1_u8 is not None: del warped_img1_u8 + if 'mask_warped' in locals(): del mask_warped + if 'mask_img2' in locals(): del mask_img2 + if 'overlap_mask' in locals(): del overlap_mask + if 'img2_part' in locals(): del img2_part # From placing img2 + + if 'output_img' in locals() and output_img is not None and output_img is not final_output_img: + # Delete intermediate output_img if it wasn't the final result (e.g., error occurred) + del output_img + + # Transformation matrices and points + if 'H_matrix_3x3_for_canvas' in locals(): del H_matrix_3x3_for_canvas + if 'final_warp_M' in locals(): del final_warp_M + if 'mask_trans' in locals(): del mask_trans + if 'src_pts' in locals(): del src_pts + if 'dst_pts' in locals(): del dst_pts + + # Feature matching intermediates + if 'kp1' in locals(): del kp1 + if 'kp2' in locals(): del kp2 + if 'des1' in locals(): del des1 + if 'des2' in locals(): del des2 + if 'good_matches' in locals(): del good_matches + if 'all_matches' in locals(): del all_matches + + # Initial uint8 images + if 'img1_u8' in locals(): del img1_u8 + if 'img2_u8' in locals(): del img2_u8 + + gc.collect() + + + return final_output_img, log_message + +# --- Function for N-Image Stitching (Primarily for Image List Input) --- +def stitch_multiple_images(images, # List of NumPy images (BGR, potentially pre-cropped) + stitcher_mode_str="SCANS", + registration_resol=0.6, + seam_estimation_resol=0.1, + compositing_resol=-1.0, # Use -1.0 for default/auto + wave_correction=False, + exposure_comp_type_str="GAIN_BLOCKS", + enable_cropping=True, # This is for POST-stitch cropping + strict_no_black_edges=False, + # Pairwise/Fallback specific params + transform_model_str="Homography", + blend_method="multi-band", + enable_gain_compensation=True, + orb_nfeatures=2000, + match_ratio_thresh=0.75, + ransac_reproj_thresh=5.0, + max_distance_coeff=0.5, + max_blending_width=10000, + max_blending_height=10000, + blend_smooth_ksize=15, + num_blend_levels=4 + ): + """ + Stitches a list of images (NumPy arrays). Tries cv2.Stitcher first (unless + stitcher_mode_str is 'DIRECT_PAIRWISE'), otherwise falls back to manual + pairwise stitching using the specified transform_model_str. + Returns ONE stitched image and log. + Input images should be in BGR format (already potentially cropped by caller). + Output is RGB. The 'enable_cropping' param here refers to final black border cropping. + """ + log = log_and_print(f"--- Starting Stitching Process for {len(images)} Provided Images ---\n", "") + total_start_time = time.time() + stitched_img_rgb = None # Initialize result + + if len(images) < 2: + log = log_and_print("Error: Need at least two images to stitch.\n", log) + return None, log + + # Check if any input image is None or empty after potential pre-cropping + valid_images = [] + for i, img in enumerate(images): + if img is None or img.size == 0: + log = log_and_print(f"Warning: Input image at index {i} is invalid (None or empty). Skipping it.\n", log) + else: + valid_images.append(img) + + if len(valid_images) < 2: + log = log_and_print(f"Error: Not enough valid images ({len(valid_images)}) left after checking. Cannot stitch.\n", log) + del images, valid_images # Clean up + gc.collect() + return None, log + + images = valid_images # Use the filtered list + log = log_and_print(f"Proceeding with {len(images)} valid images.\n", log) + log = log_and_print(f"Selected Stitcher Mode: {stitcher_mode_str}\n", log) + # Log the pairwise transform model choice, relevant if fallback or DIRECT_PAIRWISE + if stitcher_mode_str == "DIRECT_PAIRWISE": + log = log_and_print(f"Using Pairwise Transform Model: {transform_model_str}\n", log) + log = log_and_print(f"Pairwise Params: RANSAC Thresh={ransac_reproj_thresh}, Max Dist Coeff={max_distance_coeff}\n", log) + else: + log = log_and_print(f"Pairwise Transform Model (for fallback): {transform_model_str}\n", log) + log = log_and_print(f"Fallback Pairwise Params: RANSAC Thresh={ransac_reproj_thresh}, Max Dist Coeff={max_distance_coeff}\n", log) + log = log_and_print(f"Post-Crop: Enable={enable_cropping}, Strict Edges={strict_no_black_edges}\n", log) # Log new param + + + skip_cv2_stitcher = (stitcher_mode_str == "DIRECT_PAIRWISE") + + stitched_img_bgr = None + stitcher_success = False + + # 1. Try using cv2.Stitcher (unless skipped) + if not skip_cv2_stitcher: + log = log_and_print("\nAttempting stitching with built-in cv2.Stitcher...\n", log) + + # Map string parameters to OpenCV constants for cv2.Stitcher modes + stitcher_mode_map = {"PANORAMA": cv2.Stitcher_PANORAMA, "SCANS": cv2.Stitcher_SCANS} + # Default to SCANS if invalid string for cv2.Stitcher mode itself + cv2_stitcher_mode_enum = stitcher_mode_map.get(stitcher_mode_str, cv2.Stitcher_SCANS) + log = log_and_print(f"Using OpenCV Stitcher Mode Enum: {cv2_stitcher_mode_enum} (from string: {stitcher_mode_str})\n", log) + + exposure_comp_map = { + "NO": cv2.detail.ExposureCompensator_NO, + "GAIN": cv2.detail.ExposureCompensator_GAIN, + "GAIN_BLOCKS": cv2.detail.ExposureCompensator_GAIN_BLOCKS + } + exposure_comp_type = exposure_comp_map.get(exposure_comp_type_str, cv2.detail.ExposureCompensator_GAIN_BLOCKS) + log = log_and_print(f"Using Exposure Compensation: {exposure_comp_type_str}\n", log) + log = log_and_print(f"Wave Correction Enabled: {wave_correction}\n", log) + + stitcher = None # Initialize stitcher object variable + try: + stitcher = cv2.Stitcher.create(cv2_stitcher_mode_enum) + if stitcher is None: + raise RuntimeError("cv2.Stitcher.create returned None.") + + log = log_and_print(f"Setting Stitcher resolutions: Reg={registration_resol:.2f}, Seam={seam_estimation_resol:.2f}, Comp={compositing_resol:.2f}\n", log) + try: + if hasattr(stitcher, 'setRegistrationResol'): + stitcher.setRegistrationResol(float(registration_resol)) + if hasattr(stitcher, 'setSeamEstimationResol'): + stitcher.setSeamEstimationResol(float(seam_estimation_resol)) + if hasattr(stitcher, 'setCompositingResol'): + stitcher.setCompositingResol(float(compositing_resol)) + except Exception as e_res: + log = log_and_print(f"Warning: Could not set stitcher resolutions: {e_res}\n", log) + + try: + if hasattr(stitcher, 'setWaveCorrection'): + stitcher.setWaveCorrection(wave_correction) + except Exception as e_wave: + log = log_and_print(f"Warning: Could not set wave correction: {e_wave}\n", log) + + try: + if hasattr(stitcher, 'setExposureCompensator'): + compensator = cv2.detail.ExposureCompensator_createDefault(exposure_comp_type) + stitcher.setExposureCompensator(compensator) + del compensator # Release compensator object reference + except Exception as e_exp: + log = log_and_print(f"Warning: Could not set exposure compensator: {e_exp}\n", log) + + # Ensure all images are uint8 before passing to stitcher + images_uint8 = [] + for img in images: + if img.dtype != np.uint8: + images_uint8.append(img.clip(0, 255).astype(np.uint8)) + else: + images_uint8.append(img) + + status = cv2.Stitcher_ERR_NEED_MORE_IMGS # Initialize status to a known failure code + stitched_img_raw = None + + try: + log = log_and_print("Executing stitcher.stitch()...\n", log) + status, stitched_img_raw = stitcher.stitch(images_uint8) # Input 'images' should be BGR uint8 + log = log_and_print(f"stitcher.stitch() returned status: {status}\n", log) # Log the status code + + except cv2.error as e_stitch: + log = log_and_print(f"OpenCV Error occurred DURING stitcher.stitch() call: {e_stitch}\n", log) + log = log_and_print(f"Traceback:\n{traceback.format_exc()}\n", log) + log = log_and_print("Falling back to manual pairwise stitching method due to stitch() error.\n", log) + status = -99 # Set status to a custom failure code to ensure fallback + stitched_img_raw = None + except Exception as e_stitch_other: + log = log_and_print(f"Unexpected Error occurred DURING stitcher.stitch() call: {e_stitch_other}\n", log) + log = log_and_print(f"Traceback:\n{traceback.format_exc()}\n", log) + log = log_and_print("Falling back to manual pairwise stitching method due to unexpected stitch() error.\n", log) + status = -100 # Set status to a custom failure code + stitched_img_raw = None + finally: + del images_uint8 + gc.collect() + + if status == cv2.Stitcher_OK: + log = log_and_print("cv2.Stitcher successful!\n", log) + if stitched_img_raw is not None and stitched_img_raw.size > 0: + log = log_and_print(f"Stitcher output dimensions (raw): {stitched_img_raw.shape}\n", log) + # Apply FINAL black border cropping if enabled + cropped_result = crop_black_borders(stitched_img_raw, enable_cropping, strict_no_black_edges) + if cropped_result is not None and cropped_result.size > 0 : + stitched_img_bgr = cropped_result + log = log_and_print(f"Final dimensions after POST-stitch cropping: {stitched_img_bgr.shape}\n", log) + else: + stitched_img_bgr = stitched_img_raw + log = log_and_print("POST-stitch cropping failed or disabled, using raw stitcher output.\n", log) + stitcher_success = True + del stitched_img_raw + if 'cropped_result' in locals() and cropped_result is not stitched_img_bgr: + del cropped_result + gc.collect() + else: + log = log_and_print("Error: cv2.Stitcher returned status OK but the image is empty.\n", log) + else: + error_codes = { getattr(cv2, k): k for k in dir(cv2) if k.startswith('Stitcher_ERR_') } + error_codes[-99] = "ERR_STITCH_CV_ERROR" + error_codes[-100] = "ERR_STITCH_EXCEPTION" + # Check if fallback message was already logged by exceptions during stitch() + if "Falling back to manual pairwise stitching method due to" not in log.splitlines()[-5:]: + log = log_and_print(f"cv2.Stitcher failed with status code: {status} ({error_codes.get(status, f'Unknown Error {status}')})\n", log) + log = log_and_print("Falling back to manual pairwise stitching method...\n", log) + + except AttributeError as e_attr: + log = log_and_print(f"AttributeError during Stitcher setup ({e_attr}). Falling back.\n{traceback.format_exc()}\n", log) + except RuntimeError as e_runtime: + log = log_and_print(f"RuntimeError during Stitcher setup ({e_runtime}). Falling back.\n{traceback.format_exc()}\n", log) + except cv2.error as e: + log = log_and_print(f"OpenCV Error during Stitcher operation: {e}. Falling back.\n", log) + if "OutOfMemoryError" in str(e) or "Insufficient memory" in str(e): + log = log_and_print(">>> Specific OutOfMemoryError detected. Reduce resolutions or use more RAM.\n", log) + log = log_and_print(f"{traceback.format_exc()}\n", log) + except Exception as e: + log = log_and_print(f"Unexpected error during Stitcher: {e}. Falling back.\n{traceback.format_exc()}\n", log) + finally: + if stitcher is not None: + # Attempt to release stitcher resources if possible (may not exist) + try: + del stitcher + except NameError: + pass + gc.collect() + + # 2. Fallback or Direct Pairwise Stitching + # Trigger if cv2.Stitcher was skipped OR if it failed + if skip_cv2_stitcher or not stitcher_success: + # Add clearer logging based on the reason + if skip_cv2_stitcher: + log = log_and_print(f"\n--- Starting Sequential Pairwise Stitching (Direct Mode, Transform: {transform_model_str}) ---\n", log) + else: + log = log_and_print(f"\n--- Starting Sequential Pairwise Stitching (Fallback, Transform: {transform_model_str}) ---\n", log) + + if len(images) >= 2: + # Start with the first valid image. Ensure it's uint8. + if images[0].dtype != np.uint8: + current_stitched_image = images[0].clip(0, 255).astype(np.uint8) + else: + current_stitched_image = images[0].copy() # Copy to avoid modifying original list item + + sequential_stitch_success = True + for i in range(1, len(images)): + log = log_and_print(f"\nSequentially stitching image {i+1} of {len(images)} using pairwise method...\n", log) + + # Ensure next image is uint8 + if images[i].dtype != np.uint8: + next_image = images[i].clip(0, 255).astype(np.uint8) + else: + next_image = images[i] # Can use directly if already uint8 + + result, pairwise_log = stitch_pairwise_images( + current_stitched_image, # BGR uint8 + next_image, # BGR uint8 + transform_model_str=transform_model_str, + blend_method=blend_method, + enable_gain_compensation=enable_gain_compensation, + orb_nfeatures=orb_nfeatures, + match_ratio_thresh=match_ratio_thresh, + ransac_reproj_thresh=ransac_reproj_thresh, + max_distance_coeff=max_distance_coeff, + max_blending_width=max_blending_width, + max_blending_height=max_blending_height, + blend_smooth_ksize=blend_smooth_ksize, + num_blend_levels=num_blend_levels + ) + log += pairwise_log + + if result is None: + log = log_and_print(f"Error: Failed to stitch image {i+1} onto previous composite in the pairwise step. Aborting sequential process.\n", log) # Corrected index in log + sequential_stitch_success = False + if 'current_stitched_image' in locals() and current_stitched_image is not None: + del current_stitched_image # Clean up intermediate result + gc.collect() + break + + # Release the previous intermediate image before assigning the new one + if 'current_stitched_image' in locals() and current_stitched_image is not None: + del current_stitched_image + gc.collect() + current_stitched_image = result # Result is BGR uint8 + log = log_and_print(f"Intermediate stitched shape: {current_stitched_image.shape}\n", log) + # Ensure next_image is cleaned up if it was a conversion + if next_image is not images[i]: + del next_image + gc.collect() + + if sequential_stitch_success and current_stitched_image is not None: + log = log_and_print("\nSequential pairwise stitching complete. Applying final cropping...\n", log) + # Apply FINAL black border cropping if enabled + cropped_fallback = crop_black_borders(current_stitched_image, enable_cropping, strict_no_black_edges) + if cropped_fallback is not None and cropped_fallback.size > 0: + stitched_img_bgr = cropped_fallback + log = log_and_print(f"Final dimensions after POST-stitch cropping: {stitched_img_bgr.shape}\n", log) + else: + stitched_img_bgr = current_stitched_image # Use uncropped if cropping failed + log = log_and_print("POST-stitch cropping failed or disabled, using uncropped manual result.\n", log) + # Clean up the last intermediate/uncropped result if cropping was successful and created a new object + if cropped_fallback is not current_stitched_image and current_stitched_image is not None: + del current_stitched_image + if 'cropped_fallback' in locals() and cropped_fallback is not stitched_img_bgr: + del cropped_fallback + gc.collect() + else: + log = log_and_print("Sequential pairwise stitching process could not produce a final result.\n", log) + # Ensure cleanup if loop broke early or current_stitched_image was None/deleted + if 'current_stitched_image' in locals() and current_stitched_image is not None: + del current_stitched_image + gc.collect() + else: # Handle len(images) < 2 case (shouldn't happen due to initial check, but safety) + log = log_and_print("Error: Not enough images for pairwise stitching (internal check).\n", log) + + # Clean up the input image list now that it's processed + del images + if 'valid_images' in locals(): del valid_images # Should be same as images now + gc.collect() + + # 3. Final Result Check and Return + total_end_time = time.time() + log = log_and_print(f"\nTotal processing time: {total_end_time - total_start_time:.2f} seconds.\n", log) + + if stitched_img_bgr is not None and stitched_img_bgr.size > 0: + log = log_and_print("Stitching process finished for image list.", log) + try: + stitched_img_rgb = cv2.cvtColor(stitched_img_bgr, cv2.COLOR_BGR2RGB) # Convert BGR to RGB for Gradio + del stitched_img_bgr # Release BGR version + gc.collect() + return stitched_img_rgb, log + except cv2.error as e_cvt: + log = log_and_print(f"\nError converting final image to RGB: {e_cvt}. Returning None.\n", log) + if 'stitched_img_bgr' in locals(): del stitched_img_bgr + gc.collect() + return None, log + else: + log = log_and_print("Error: Stitching failed. No final image generated.", log) + if 'stitched_img_bgr' in locals() and stitched_img_bgr is not None: + del stitched_img_bgr + gc.collect() + return None, log + + +# --- Video Frame Stitching --- +def stitch_video_frames(video_path, + crop_top_percent=0.0, + crop_bottom_percent=0.0, + enable_cropping=True, # This is for POST-stitch cropping + strict_no_black_edges=False, + # Pairwise specific params + transform_model_str="Homography", + blend_method="multi-band", + enable_gain_compensation=True, + orb_nfeatures=2000, + match_ratio_thresh=0.75, + ransac_reproj_thresh=5.0, + max_distance_coeff=0.5, + max_blending_width=10000, + max_blending_height=10000, + blend_smooth_ksize=15, + num_blend_levels=4, + # Video specific params + sample_interval_ms=3000, + max_composite_width=10000, + max_composite_height=10000, + progress=None): + """ + Reads a video, samples frames incrementally, applies percentage crop, + and stitches them sequentially using the specified transform_model_str. + Includes size checks to limit composite image growth. + Returns a list of stitched images (RGB format) and a log. + """ + log = log_and_print(f"--- Starting Incremental Video Stitching for: {os.path.basename(video_path)} ---\n", "") + log = log_and_print(f"Params: Interval={sample_interval_ms}ms, Transform={transform_model_str}, ORB={orb_nfeatures}, Ratio={match_ratio_thresh}\n", log) + log = log_and_print(f"Params Cont'd: RANSAC Thresh={ransac_reproj_thresh}, Max Dist Coeff={max_distance_coeff}\n", log) + log = log_and_print(f"Composite Limits: MaxW={max_composite_width}, MaxH={max_composite_height}\n", log) + log = log_and_print(f"Pre-Crop: Top={crop_top_percent}%, Bottom={crop_bottom_percent}%\n", log) + log = log_and_print(f"Post-Crop Black Borders: {enable_cropping}, Strict Edges: {strict_no_black_edges}\n", log) + log = log_and_print(f"Blending: Method={blend_method}, GainComp={enable_gain_compensation}, SmoothKSize={blend_smooth_ksize}, MB Levels={num_blend_levels}\n", log) + total_start_time = time.time() + stitched_results_rgb = [] # Store final RGB images + + cap = cv2.VideoCapture(video_path) + if not cap.isOpened(): + log = log_and_print(f"Error: Could not open video file: {video_path}\n", log) + return [], log + + fps = cap.get(cv2.CAP_PROP_FPS) + frame_count_total = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) + if fps <= 0 or np.isnan(fps): # Handle invalid FPS reads + fps = 30 # Default FPS + log = log_and_print("Warning: Could not read valid FPS, defaulting to 30.\n", log) + if frame_count_total <= 0: # Handle invalid frame count reads + log = log_and_print("Warning: Could not read valid total frame count.\n", log) + total_sampled_estimate = 0 # Cannot estimate progress accurately + else: + # Estimate total frames to be sampled, avoid division by zero if interval is 0 + frames_per_sample = max(1, int(round(fps * (sample_interval_ms / 1000.0)))) if sample_interval_ms > 0 else frame_count_total + total_sampled_estimate = frame_count_total / frames_per_sample if frames_per_sample > 0 else 0 + + + frame_interval = max(1, int(round(fps * (sample_interval_ms / 1000.0)))) + log = log_and_print(f"Video Info: ~{fps:.2f} FPS, {frame_count_total} Frames, Sampling every {frame_interval} frames.\n", log) + + frame_num = 0 + processed_sampled_count = 0 # Counter for progress bar + anchor_frame = None # The starting frame of the current sequence (BGR, cropped) + current_composite = None # The stitched result being built (BGR, uint8) + last_saved_composite = None # Keep track of the last saved image to avoid duplicates + + while True: + frame_bgr_raw = None # Initialize here for cleanup later + try: + if cap is None or not cap.isOpened(): + log = log_and_print("Error: Video capture became invalid during processing.\n", log) + break + ret, frame_bgr_raw = cap.read() + if not ret: + log = log_and_print("\nEnd of video stream reached.\n", log) + break # End of video + + # --- Sampling Logic --- + if frame_num % frame_interval == 0: + if frame_bgr_raw is not None and frame_bgr_raw.size > 0: + processed_sampled_count += 1 + frame_bgr = None # Initialize BGR frame variable + + # --- Frame Format Check --- + if frame_bgr_raw.ndim == 2: + frame_bgr = cv2.cvtColor(frame_bgr_raw, cv2.COLOR_GRAY2BGR) + elif frame_bgr_raw.ndim == 3 and frame_bgr_raw.shape[2] == 4: + frame_bgr = cv2.cvtColor(frame_bgr_raw, cv2.COLOR_BGRA2BGR) + elif frame_bgr_raw.ndim == 3 and frame_bgr_raw.shape[2] == 3: + frame_bgr = frame_bgr_raw # Already BGR + else: + log = log_and_print(f"Warning: Skipping frame {frame_num} due to unexpected shape {frame_bgr_raw.shape}\n", log) + if frame_bgr_raw is not None: del frame_bgr_raw # Clean up original frame + gc.collect() + frame_num += 1 + continue # Skip to next frame read + + # Release the raw frame once converted/checked (if a copy was made) + if frame_bgr is not frame_bgr_raw: + del frame_bgr_raw + frame_bgr_raw = None # Mark as deleted + gc.collect() + + cropped_frame_bgr = crop_image_by_percent(frame_bgr, crop_top_percent, crop_bottom_percent) + del frame_bgr # Release the uncropped BGR version + gc.collect() + + # Check if cropping failed or resulted in an empty image + if cropped_frame_bgr is None or cropped_frame_bgr.size == 0: + log = log_and_print(f"Warning: Skipping frame {frame_num} because percentage cropping failed or resulted in empty image.\n", log) + if cropped_frame_bgr is not None: del cropped_frame_bgr # Should be None, but safety check + gc.collect() + frame_num += 1 + continue # Skip to next frame read + + # Now use 'cropped_frame_bgr' as the current frame for stitching + current_frame_for_stitch = cropped_frame_bgr # BGR, uint8, potentially cropped + + if progress is not None and total_sampled_estimate > 0: + # Ensure progress doesn't exceed 1.0 + progress_fraction = min(1.0, processed_sampled_count / total_sampled_estimate) + progress(progress_fraction, desc=f"Processing Sample {processed_sampled_count}/{int(total_sampled_estimate)}") + elif progress is not None: + # Fallback progress if estimate is bad + progress(processed_sampled_count / (processed_sampled_count + 10), desc=f"Processing Sample {processed_sampled_count}") + + + log = log_and_print(f"\n--- Processing sampled frame index {frame_num} (Count: {processed_sampled_count}) ---\n", log) + log = log_and_print(f"Frame shape after potential pre-crop: {current_frame_for_stitch.shape}\n", log) + + # --- Stitching Logic --- + if anchor_frame is None: + # Start a new sequence + anchor_frame = current_frame_for_stitch.copy() # Make a copy + current_composite = anchor_frame # Start composite is the anchor itself + log = log_and_print(f"Frame {frame_num}: Set as new anchor (Shape: {anchor_frame.shape}).\n", log) + # No need to stitch yet, just set the anchor + + else: + # Try stitching the current composite with the new frame + log = log_and_print(f"Attempting stitch: Composite({current_composite.shape}) + Frame({current_frame_for_stitch.shape})\n", log) + + stitch_result, stitch_log = stitch_pairwise_images( + current_composite, # Previous result or anchor (uint8) + current_frame_for_stitch, # New frame to add (uint8) + transform_model_str=transform_model_str, + blend_method=blend_method, + enable_gain_compensation=enable_gain_compensation, + orb_nfeatures=orb_nfeatures, + match_ratio_thresh=match_ratio_thresh, + ransac_reproj_thresh=ransac_reproj_thresh, + max_distance_coeff=max_distance_coeff, + max_blending_width=max_blending_width, + max_blending_height=max_blending_height, + blend_smooth_ksize=blend_smooth_ksize, + num_blend_levels=num_blend_levels + ) + log += stitch_log + + if stitch_result is not None and stitch_result.size > 0: + # --- Stitching SUCCEEDED --- + log = log_and_print(f"Success: Stitched frame {frame_num}. New composite shape: {stitch_result.shape}\n", log) + # Release old composite before assigning new one + del current_composite + gc.collect() + current_composite = stitch_result # Update the composite (stitch_result is BGR uint8) + # anchor_frame remains the same for this sequence + + # --- Check Size Limit --- + h_curr, w_curr = current_composite.shape[:2] + size_limit_exceeded = False + # Check only if limit > 0 + if max_composite_width > 0 and w_curr > max_composite_width: + log = log_and_print(f"ACTION: Composite width ({w_curr}) exceeded limit ({max_composite_width}).\n", log) + size_limit_exceeded = True + if max_composite_height > 0 and h_curr > max_composite_height: + log = log_and_print(f"ACTION: Composite height ({h_curr}) exceeded limit ({max_composite_height}).\n", log) + size_limit_exceeded = True + + if size_limit_exceeded: + log = log_and_print("Saving current composite and starting new sequence with NEXT frame.\n", log) + + # Apply FINAL black border cropping if enabled + post_cropped_composite = crop_black_borders(current_composite, enable_cropping, strict_no_black_edges) + if post_cropped_composite is not None and post_cropped_composite.size > 0: + # Avoid saving the exact same image twice in a row + is_duplicate = False + if last_saved_composite is not None: + try: + # Simple check: compare shapes first, then content if shapes match + if last_saved_composite.shape == post_cropped_composite.shape: + if np.array_equal(last_saved_composite, post_cropped_composite): + is_duplicate = True + except Exception as e_comp: + log = log_and_print(f"Warning: Error comparing images for duplication check: {e_comp}\n", log) + + if not is_duplicate: + try: + stitched_results_rgb.append(cv2.cvtColor(post_cropped_composite, cv2.COLOR_BGR2RGB)) + # Update last_saved_composite only if append is successful + if last_saved_composite is not None: del last_saved_composite + last_saved_composite = post_cropped_composite.copy() # Store the saved one (BGR) + log = log_and_print(f"Saved composite image {len(stitched_results_rgb)} (Post-Cropped Shape: {post_cropped_composite.shape}).\n", log) + except cv2.error as e_cvt: + log = log_and_print(f"Warning: Failed to convert size-limited composite to RGB: {e_cvt}\n", log) + except Exception as e_save: + log = log_and_print(f"Warning: Failed to save size-limited composite: {e_save}\n", log) + + else: + log = log_and_print("Skipping save: Result identical to previously saved image.\n", log) + + # Clean up the post-cropped version if it wasn't stored in last_saved_composite + if last_saved_composite is not post_cropped_composite: + del post_cropped_composite + gc.collect() + else: + log = log_and_print("Warning: Post-stitch cropping failed for the size-limited composite, skipping save.\n", log) + if post_cropped_composite is not None: del post_cropped_composite # Delete if it existed but was empty + + # Reset for the next frame to become the anchor + del current_composite + if anchor_frame is not None: del anchor_frame # Delete old anchor too + if last_saved_composite is not None: del last_saved_composite # Reset duplicate check too + current_composite = None + anchor_frame = None + last_saved_composite = None + gc.collect() + # --- End Size Check --- + + else: + # --- Stitching FAILED --- + log = log_and_print(f"Failed: Could not stitch frame {frame_num} onto current composite.\n", log) + # Save the *previous* valid composite (if it exists and is not just the anchor) + save_previous = False + if current_composite is not None and anchor_frame is not None: + # Check if composite is actually different from the anchor + try: + if current_composite.shape != anchor_frame.shape or not np.array_equal(current_composite, anchor_frame): + save_previous = True + except Exception as e_comp: + log = log_and_print(f"Warning: Error comparing composite to anchor: {e_comp}\n", log) + save_previous = True # Assume different if compare fails + + if save_previous: + log = log_and_print("ACTION: Saving the previously stitched result before resetting.\n", log) + # Apply FINAL black border cropping if enabled + post_cropped_composite = crop_black_borders(current_composite, enable_cropping, strict_no_black_edges) + if post_cropped_composite is not None and post_cropped_composite.size > 0: + is_duplicate = False + if last_saved_composite is not None: + try: + if last_saved_composite.shape == post_cropped_composite.shape: + if np.array_equal(last_saved_composite, post_cropped_composite): + is_duplicate = True + except Exception as e_comp: + log = log_and_print(f"Warning: Error comparing images for duplication check: {e_comp}\n", log) + + if not is_duplicate: + try: + stitched_results_rgb.append(cv2.cvtColor(post_cropped_composite, cv2.COLOR_BGR2RGB)) + if last_saved_composite is not None: del last_saved_composite + last_saved_composite = post_cropped_composite.copy() # Store BGR + log = log_and_print(f"Saved composite image {len(stitched_results_rgb)} (Post-Cropped Shape: {post_cropped_composite.shape}).\n", log) + except cv2.error as e_cvt: + log = log_and_print(f"Warning: Failed to convert previous composite to RGB: {e_cvt}\n", log) + except Exception as e_save: + log = log_and_print(f"Warning: Failed to save previous composite: {e_save}\n", log) + else: + log = log_and_print("Skipping save: Result identical to previously saved image.\n", log) + + if last_saved_composite is not post_cropped_composite: + del post_cropped_composite + gc.collect() + else: + log = log_and_print("Warning: Post-stitch cropping failed for the previously stitched result, skipping save.\n", log) + if post_cropped_composite is not None: del post_cropped_composite + else: + log = log_and_print("No previous composite to save (stitching failed on first attempt for this anchor or composite was just the anchor).\n", log) + + # The frame that *failed* to stitch becomes the new anchor + log = log_and_print(f"ACTION: Setting frame {frame_num} (shape: {current_frame_for_stitch.shape}) as the new anchor.\n", log) + if current_composite is not None: del current_composite # Delete the old composite + if anchor_frame is not None: del anchor_frame # Delete the old anchor + if last_saved_composite is not None: del last_saved_composite # Reset duplicate check + gc.collect() + anchor_frame = current_frame_for_stitch.copy() # Use the frame that failed (already cropped) + current_composite = anchor_frame # Reset composite to this new anchor + last_saved_composite = None + gc.collect() + # current_frame_for_stitch is now anchor_frame, no need to delete separately below + + # --- Clean up current frame AFTER processing (if it wasn't made the new anchor) --- + # If stitching succeeded OR if it failed but wasn't the first frame, + # current_frame_for_stitch needs cleanup unless it just became the anchor. + if 'current_frame_for_stitch' in locals() and current_frame_for_stitch is not anchor_frame: + del current_frame_for_stitch + gc.collect() + + else: # Handle cases where frame_bgr_raw is None or empty after read + if frame_bgr_raw is not None: + del frame_bgr_raw + frame_bgr_raw = None + gc.collect() + else: # Frame not sampled + # Still need to release the raw frame if it was read + if frame_bgr_raw is not None: + del frame_bgr_raw + frame_bgr_raw = None + # Don't gc.collect() on every skipped frame, too slow + + frame_num += 1 + # Loop continues + except Exception as loop_error: + log = log_and_print(f"Unexpected error in main video loop at frame {frame_num}: {loop_error}\n{traceback.format_exc()}\n", log) + # Try to continue to next frame if possible, or break if capture seems broken + if cap is None or not cap.isOpened(): + log = log_and_print("Video capture likely broken, stopping loop.\n", log) + break + else: + frame_num += 1 # Ensure frame counter increments + # Clean up potentially lingering frame data from the failed iteration + if 'frame_bgr_raw' in locals() and frame_bgr_raw is not None: del frame_bgr_raw + if 'frame_bgr' in locals() and frame_bgr is not None: del frame_bgr + if 'cropped_frame_bgr' in locals() and cropped_frame_bgr is not None: del cropped_frame_bgr + if 'current_frame_for_stitch' in locals() and current_frame_for_stitch is not None and current_frame_for_stitch is not anchor_frame: del current_frame_for_stitch + gc.collect() + + # --- After the loop: Check if there's a final composite to save --- + if current_composite is not None and anchor_frame is not None: + # Only save if it contains more than just the last anchor frame OR if it's the *only* result + save_final = False + if len(stitched_results_rgb) == 0: # If no images saved yet, save this one + save_final = True + else: + try: + if current_composite.shape != anchor_frame.shape or not np.array_equal(current_composite, anchor_frame): + save_final = True + except Exception as e_comp: + log = log_and_print(f"Warning: Error comparing final composite to anchor: {e_comp}\n", log) + save_final = True # Save if comparison fails + + if save_final: + log = log_and_print("\nEnd of frames reached. Checking final composite...\n", log) + post_cropped_final = crop_black_borders(current_composite, enable_cropping, strict_no_black_edges) + if post_cropped_final is not None and post_cropped_final.size > 0: + is_duplicate = False + if last_saved_composite is not None: + try: + if last_saved_composite.shape == post_cropped_final.shape: + if np.array_equal(last_saved_composite, post_cropped_final): + is_duplicate = True + except Exception as e_comp: + log = log_and_print(f"Warning: Error comparing final image for duplication check: {e_comp}\n", log) + + if not is_duplicate: + try: + stitched_results_rgb.append(cv2.cvtColor(post_cropped_final, cv2.COLOR_BGR2RGB)) + log = log_and_print(f"Saved final composite image {len(stitched_results_rgb)} (Post-Cropped Shape: {post_cropped_final.shape}).\n", log) + # No need to update last_saved_composite here, loop is finished + except cv2.error as e_cvt: + log = log_and_print(f"Warning: Failed to convert final composite to RGB: {e_cvt}\n", log) + except Exception as e_save: + log = log_and_print(f"Warning: Failed to save final composite: {e_save}\n", log) + else: + log = log_and_print("Skipping save of final composite: Result identical to previously saved image.\n", log) + + # Clean up final cropped image if it existed + del post_cropped_final + gc.collect() + else: + log = log_and_print("Warning: Post-stitch cropping failed for the final composite, skipping save.\n", log) + if post_cropped_final is not None: del post_cropped_final # Delete if empty + else: + log = log_and_print("\nEnd of frames reached. Final composite was identical to its anchor frame and not the only result, not saving.\n", log) + + # --- Final Cleanup --- + if cap is not None and cap.isOpened(): + cap.release() + if 'cap' in locals(): del cap + if 'anchor_frame' in locals() and anchor_frame is not None: del anchor_frame + if 'current_composite' in locals() and current_composite is not None: del current_composite + if 'last_saved_composite' in locals() and last_saved_composite is not None: del last_saved_composite + gc.collect() + + total_end_time = time.time() + log = log_and_print(f"\nVideo stitching process finished. Found {len(stitched_results_rgb)} stitched image(s).", log) + log = log_and_print(f"\nTotal processing time: {total_end_time - total_start_time:.2f} seconds.\n", log) + + # Filter out potential None entries just before returning + final_results = [img for img in stitched_results_rgb if img is not None and img.size > 0] + if len(final_results) != len(stitched_results_rgb): + log = log_and_print(f"Warning: Filtered out {len(stitched_results_rgb) - len(final_results)} None or empty results before final return.\n", log) + # Clean up the original list with potential Nones + del stitched_results_rgb + gc.collect() + + return final_results, log + + +# --- Gradio Interface Function --- +def run_stitching_interface(input_files, + crop_top_percent, + crop_bottom_percent, + stitcher_mode_str, # For cv2.Stitcher + registration_resol, + seam_estimation_resol, + compositing_resol, + wave_correction, + exposure_comp_type_str, # For cv2.Stitcher + enable_cropping, # Post-stitch black border crop + strict_no_black_edges_input, + # Detailed Stitcher Settings + transform_model_str, + blend_method_str, + enable_gain_compensation, + orb_nfeatures, + match_ratio_thresh, + ransac_reproj_thresh_input, + max_distance_coeff_input, + max_blending_width, + max_blending_height, + blend_smooth_ksize_input, + num_blend_levels_input, + # Video specific settings + sample_interval_ms, + max_composite_width_video, + max_composite_height_video, + progress=gr.Progress(track_tqdm=True) + ): + """ + Wrapper function called by the Gradio interface. + Handles input (images or video), applies pre-cropping, + calls the appropriate stitching logic (passing transform_model_str), + and returns results. + """ + if input_files is None or len(input_files) == 0: + return [], "Please upload images or a video file." + + # Convert Gradio inputs to correct types + blend_smooth_ksize = int(blend_smooth_ksize_input) if blend_smooth_ksize_input is not None else -1 + num_blend_levels = int(num_blend_levels_input) if num_blend_levels_input is not None else 4 + ransac_reproj_thresh = float(ransac_reproj_thresh_input) if ransac_reproj_thresh_input is not None else 3.0 + max_distance_coeff = float(max_distance_coeff_input) if max_distance_coeff_input is not None else 0.5 + + log = f"Received {len(input_files)} file(s).\n" + log = log_and_print(f"Pre-Crop Settings: Top={crop_top_percent}%, Bottom={crop_bottom_percent}%\n", log) + log = log_and_print(f"Post-Crop Black Borders: Enabled={enable_cropping}, Strict Edges={strict_no_black_edges_input}\n", log) + # Log detailed settings including new ones + log = log_and_print(f"Detailed Settings: Transform={transform_model_str}, Blend={blend_method_str}, GainComp={enable_gain_compensation}, ORB={orb_nfeatures}, Ratio={match_ratio_thresh}\n", log) + log = log_and_print(f"Detailed Settings Cont'd: RANSAC Thresh={ransac_reproj_thresh}, MaxDistCoeff={max_distance_coeff}, MaxBlendW={max_blending_width}, MaxBlendH={max_blending_height}, SmoothKSize={blend_smooth_ksize}, MBLevels={num_blend_levels}\n", log) + progress(0, desc="Processing Input...") + + # Determine input type: List of images or a single video + is_video_input = False + video_path = None + image_paths = [] + + # Check file types using mimetypes + try: + # Handle potential TempfileWrappers or string paths + input_filepaths = [] + for f in input_files: + if hasattr(f, 'name'): # Gradio File object + input_filepaths.append(f.name) + elif isinstance(f, str): # String path (e.g., from examples) + input_filepaths.append(f) + else: + log = log_and_print(f"Warning: Unexpected input file type: {type(f)}. Skipping.\n", log) + + + if len(input_filepaths) == 1: + filepath = input_filepaths[0] + mime_type, _ = mimetypes.guess_type(filepath) + if mime_type and mime_type.startswith('video'): + is_video_input = True + video_path = filepath + log = log_and_print(f"Detected video input: {os.path.basename(video_path)}\n", log) + elif mime_type and mime_type.startswith('image'): + log = log_and_print("Detected single image input. Need at least two images for list stitching.\n", log) + image_paths = [filepath] # Keep it for error message later + else: + # Fallback check: try reading as image + img_test = None + try: + # Use np.fromfile for paths that might have unicode characters + n = np.fromfile(filepath, np.uint8) + if n.size > 0: + img_test = cv2.imdecode(n, cv2.IMREAD_COLOR) + else: + raise ValueError("File is empty") + + if img_test is not None and img_test.size > 0: + log = log_and_print(f"Warning: Unknown file type for single file: {os.path.basename(filepath)}. Assuming image based on successful read. Need >= 2 images.\n", log) + image_paths = [filepath] + del img_test + else: + raise ValueError("Cannot read as image or image is empty") + except Exception as e_read_check: + log = log_and_print(f"Error: Could not determine file type or read single file: {os.path.basename(filepath)}. Error: {e_read_check}. Please provide video or image files.\n", log) + if img_test is not None: del img_test + return [], log + else: # Multiple files uploaded + image_paths = [] + non_image_skipped = False + for filepath in input_filepaths: + mime_type, _ = mimetypes.guess_type(filepath) + is_image = False + if mime_type and mime_type.startswith('image'): + is_image = True + else: + # Fallback check: Try reading as image + img_test = None + try: + n = np.fromfile(filepath, np.uint8) + if n.size > 0: + img_test = cv2.imdecode(n, cv2.IMREAD_COLOR) + else: + raise ValueError("File is empty") + + if img_test is not None and img_test.size > 0: + is_image = True + log = log_and_print(f"Warning: Non-image or unknown file type detected: {os.path.basename(filepath)}. Assuming image based on read success.\n", log) + del img_test + else: + non_image_skipped = True + log = log_and_print(f"Warning: Skipping non-image file (or empty/failed read): {os.path.basename(filepath)}\n", log) + except Exception as e_read_check: + non_image_skipped = True + log = log_and_print(f"Warning: Skipping non-image file (read failed: {e_read_check}): {os.path.basename(filepath)}\n", log) + if img_test is not None: del img_test + + + if is_image: + image_paths.append(filepath) + + if not image_paths: # No valid images found + if non_image_skipped: + log = log_and_print("Error: No valid image files found in the input list after filtering.\n", log) + else: # Should only happen if initial list was empty, but covered by check at start + log = log_and_print("Error: No image files provided in the input list.\n", log) + return [], log + elif non_image_skipped: + log = log_and_print(f"Proceeding with {len(image_paths)} assumed image files (some non-images were skipped).\n", log) + else: + log = log_and_print(f"Detected {len(image_paths)} image inputs.\n", log) + except Exception as e: + log = log_and_print(f"Error detecting input file types: {e}\n{traceback.format_exc()}\n", log) + return [], log + + + # --- Process Based on Input Type --- + final_stitched_images_rgb = [] # List to hold RGB results for gallery + stitch_log = "" + + if is_video_input: + # --- VIDEO PROCESSING --- + log = log_and_print("Starting incremental video frame stitching...\n", log) + progress(0.1, desc="Sampling & Stitching Video...") + # Ensure blend method string is lowercase for internal checks + blend_method_lower = blend_method_str.lower() if blend_method_str else "multi-band" + + final_stitched_images_rgb, stitch_log = stitch_video_frames( + video_path, + crop_top_percent=crop_top_percent, + crop_bottom_percent=crop_bottom_percent, + enable_cropping=enable_cropping, # Post-stitch crop + strict_no_black_edges=strict_no_black_edges_input, + transform_model_str=transform_model_str, + blend_method=blend_method_lower, # linear or multi-band + enable_gain_compensation=enable_gain_compensation, + orb_nfeatures=orb_nfeatures, + match_ratio_thresh=match_ratio_thresh, + ransac_reproj_thresh=ransac_reproj_thresh, + max_distance_coeff=max_distance_coeff, + max_blending_width=max_blending_width, + max_blending_height=max_blending_height, + sample_interval_ms=sample_interval_ms, + max_composite_width=max_composite_width_video, + max_composite_height=max_composite_height_video, + blend_smooth_ksize=blend_smooth_ksize, + num_blend_levels=num_blend_levels, + progress=progress + ) + + elif len(image_paths) >= 2: + # --- IMAGE LIST PROCESSING --- + log = log_and_print("Reading and preparing images for list stitching...\n", log) + images_bgr_cropped = [] # Store potentially cropped BGR images + read_success = True + for i, img_path in enumerate(image_paths): + progress(i / len(image_paths) * 0.1, desc=f"Reading image {i+1}/{len(image_paths)}") # Small progress for reading + img = None + try: + n = np.fromfile(img_path, np.uint8) + if n.size > 0: + img = cv2.imdecode(n, cv2.IMREAD_UNCHANGED) + else: + log = log_and_print(f"Error: File is empty: {os.path.basename(img_path)}. Skipping.\n", log) + continue + if img is None: + raise ValueError("imdecode returned None") + except Exception as e_read: + log = log_and_print(f"Error reading image: {os.path.basename(img_path)}. Error: {e_read}. Skipping.\n", log) + if img is not None: del img + continue # Skip to the next image + + # Convert to BGR + img_bgr = None + try: + if img.ndim == 2: + img_bgr = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR) + elif img.ndim == 3 and img.shape[2] == 4: + img_bgr = cv2.cvtColor(img, cv2.COLOR_BGRA2BGR) + elif img.ndim == 3 and img.shape[2] == 3: + img_bgr = img # Already BGR, no copy needed yet + else: + log = log_and_print(f"Error: Invalid image shape {img.shape} for {os.path.basename(img_path)}. Skipping.\n", log) + del img + if 'img_bgr' in locals() and img_bgr is not None: del img_bgr + gc.collect() + continue # Skip to the next image + except cv2.error as e_cvt_color: + log = log_and_print(f"Error converting image color for {os.path.basename(img_path)}: {e_cvt_color}. Skipping.\n", log) + del img + if 'img_bgr' in locals() and img_bgr is not None: del img_bgr + gc.collect() + continue + + # Release original read image if conversion happened + if img_bgr is not img: + del img + gc.collect() + + # Apply Percentage Cropping + img_bgr_cropped_single = crop_image_by_percent(img_bgr, crop_top_percent, crop_bottom_percent) + + # Release uncropped BGR version (unless it was the result of cropping) + if img_bgr_cropped_single is not img_bgr: + del img_bgr + gc.collect() + + if img_bgr_cropped_single is None or img_bgr_cropped_single.size == 0: + log = log_and_print(f"Warning: Skipping image {os.path.basename(img_path)} because percentage cropping failed or resulted in empty image.\n", log) + if img_bgr_cropped_single is not None: del img_bgr_cropped_single + gc.collect() + continue # Skip to next image + + images_bgr_cropped.append(img_bgr_cropped_single) + # log = log_and_print(f"Read and pre-cropped: {os.path.basename(img_path)} -> Shape: {img_bgr_cropped_single.shape}\n", log) # Can be verbose + + if len(images_bgr_cropped) < 2: + stitch_log = log_and_print(f"Need at least two valid images after reading and pre-cropping ({len(images_bgr_cropped)} found) for list stitching.\n", log) # Append to main log + read_success = False # Indicate failure to proceed + else: + log = log_and_print(f"Proceeding with {len(images_bgr_cropped)} valid, pre-cropped images. Starting list stitching...\n", log) + progress(0.1, desc="Stitching Image List...") + # Ensure blend method string is lowercase for internal checks + blend_method_lower = blend_method_str.lower() if blend_method_str else "multi-band" + + stitched_single_rgb, stitch_log_img = stitch_multiple_images( + images_bgr_cropped, # Pass the list of cropped images + stitcher_mode_str=stitcher_mode_str, + registration_resol=registration_resol, + seam_estimation_resol=seam_estimation_resol, + compositing_resol=compositing_resol, + wave_correction=wave_correction, + exposure_comp_type_str=exposure_comp_type_str, + enable_cropping=enable_cropping, # Post-stitch crop + strict_no_black_edges=strict_no_black_edges_input, + transform_model_str=transform_model_str, + blend_method=blend_method_lower, + enable_gain_compensation=enable_gain_compensation, + orb_nfeatures=orb_nfeatures, + match_ratio_thresh=match_ratio_thresh, + ransac_reproj_thresh=ransac_reproj_thresh, + max_distance_coeff=max_distance_coeff, + max_blending_width=max_blending_width, + max_blending_height=max_blending_height, + blend_smooth_ksize=blend_smooth_ksize, + num_blend_levels=num_blend_levels + ) + stitch_log += stitch_log_img # Append log from stitching function + if stitched_single_rgb is not None: + final_stitched_images_rgb = [stitched_single_rgb] # Result is a list containing the single image + + # Clean up loaded images for list mode after stitching attempt + if 'images_bgr_cropped' in locals(): + for img_del in images_bgr_cropped: + if img_del is not None: del img_del + del images_bgr_cropped + gc.collect() + + elif len(image_paths) == 1: + # This case should have been handled by the input type detection, + # but add a message here just in case. + log = log_and_print("Error: Only one image file provided or detected. Need at least two for image list stitching.\n", log) + stitch_log = "" # No stitching attempted + else: + # This case means no valid input files were found or passed initial checks. + log = log_and_print("Error: Input must be a single video file or at least two image files. No valid input found.\n", log) + stitch_log = "" + + final_log = log + stitch_log + if not final_stitched_images_rgb: + # Avoid duplicating error messages if log already indicates failure + if "Error:" not in final_log[-200:]: # Check last few lines for errors + final_log = log_and_print("\nNo stitched images were generated.", final_log) + + # --- Saving Results to Temporary Files --- + output_file_paths = [] # List to store paths for the Gallery + temp_dir = None + + if final_stitched_images_rgb: + try: + # Try to create a subdirectory within the default Gradio temp space if possible + gradio_temp_base = tempfile.gettempdir() + gradio_subdir = os.path.join(gradio_temp_base, 'gradio') # Default Gradio temp subdir name + # Check if we can write there, otherwise use default temp dir + target_temp_dir_base = gradio_subdir if os.path.exists(gradio_subdir) and os.access(gradio_subdir, os.W_OK) else gradio_temp_base + + if not os.path.exists(target_temp_dir_base): + try: + os.makedirs(target_temp_dir_base) + except OSError as e_mkdir: + final_log = log_and_print(f"Warning: Could not create temp directory '{target_temp_dir_base}', using default. Error: {e_mkdir}\n", final_log) + target_temp_dir_base = tempfile.gettempdir() # Fallback to system default temp + + temp_dir = tempfile.mkdtemp(prefix="stitch_run_", dir=target_temp_dir_base) + final_log = log_and_print(f"\nInfo: Saving output images to temporary directory: {temp_dir}\n", final_log) + + for i, img_rgb in enumerate(final_stitched_images_rgb): + if img_rgb is None or img_rgb.size == 0: + final_log = log_and_print(f"Warning: Skipping saving image index {i} because it is None or empty.\n", final_log) + continue + filename = f"stitched_image_{i+1:03d}.png" + # Use os.path.join for cross-platform compatibility + full_path = os.path.join(temp_dir, filename) + img_bgr = None # Initialize for finally block + try: + img_bgr = cv2.cvtColor(img_rgb, cv2.COLOR_RGB2BGR) + # Use imencode -> write pattern for better handling of paths/special chars + is_success, buf = cv2.imencode('.png', img_bgr) + if is_success: + with open(full_path, 'wb') as f: + f.write(buf) + # Use Gradio File obj or just path string? Gallery seems to prefer path strings. + output_file_paths.append((full_path, filename)) # Store the full path for Gradio Gallery + # final_log = log_and_print(f"Successfully saved: {filename}\n", final_log) # Can be verbose + else: + final_log = log_and_print(f"Warning: Failed to encode image for saving: {filename}\n", final_log) + except cv2.error as e_cvt_write: + final_log = log_and_print(f"Error converting or encoding image {filename}: {e_cvt_write}\n", final_log) + except IOError as e_io: + final_log = log_and_print(f"Error writing image file {filename} to {full_path}: {e_io}\n", final_log) + except Exception as e_write: + final_log = log_and_print(f"Unexpected error writing image {filename} to {full_path}: {e_write}\n", final_log) + finally: + if img_bgr is not None: del img_bgr + gc.collect() + except Exception as e_tempdir: + final_log = log_and_print(f"Error creating temporary directory or saving output: {e_tempdir}\n", final_log) + output_file_paths = [] # Fallback to empty list + + # --- Final Cleanup of RGB images list --- + if 'final_stitched_images_rgb' in locals(): + for img_del in final_stitched_images_rgb: + if img_del is not None: del img_del + del final_stitched_images_rgb + gc.collect() + + progress(1.0, desc="Finished!") + final_log = log_and_print("\nCleanup complete.", final_log) + + # Return the LIST OF FILE PATHS for the Gallery, and the log + return output_file_paths, final_log + + +# --- Define Gradio Interface --- +with gr.Blocks() as demo: + gr.Markdown("# OpenCV Image and Video Stitcher") + gr.Markdown( + "Upload multiple images (for list/panorama stitching) OR a single video file (for sequential frame stitching). " + "Video frames are sampled incrementally based on the interval. " + "Use Pre-Cropping to remove unwanted areas *before* stitching. Adjust other parameters and click 'Stitch'." + ) + + with gr.Row(): + with gr.Column(scale=1): + stitch_button = gr.Button("Stitch", variant="primary") + input_files = gr.File( + label="Upload Images or a Video", + # Common image and video types + file_types=["image", ".mp4", ".avi", ".mov", ".mkv", ".wmv", ".webm"], + file_count="multiple", + elem_id="input_files" + ) + + # --- Parameters Grouping --- + with gr.Accordion("Preprocessing Settings", open=True): + crop_top_percent = gr.Slider(0.0, 49.0, step=0.5, value=0.0, label="Crop Top %", + info="Percentage of height to remove from the TOP of each image/frame BEFORE stitching.") + crop_bottom_percent = gr.Slider(0.0, 49.0, step=0.5, value=0.0, label="Crop Bottom %", + info="Percentage of height to remove from the BOTTOM of each image/frame BEFORE stitching.") + + with gr.Accordion("OpenCV Stitcher Settings (Image List Mode Only)", open=True): + stitcher_mode = gr.Radio(["SCANS", "PANORAMA", "DIRECT_PAIRWISE"], label="Stitcher Mode (Image List)", value="SCANS", + info=( + "Method for image list stitching. 'SCANS'/'PANORAMA': Use OpenCV's built-in Stitcher (optimized for translation/rotation). " + "'SCANS': Optimized for images primarily related by translation (like scanning documents or linear camera motion), potentially using simpler geometric models or assumptions internally. " + "'PANORAMA': Designed for images captured by rotating the camera around a central point. It uses full perspective transformations (Homography) to handle the complex geometric distortions typical in panoramic shots." + "'DIRECT_PAIRWISE': Skip OpenCV Stitcher and directly use sequential pairwise feature matching (same as video mode or fallback)." + ) + ) + registration_resol = gr.Slider(0.1, 1.0, step=0.05, value=0.6, label="Registration Resolution", + info="Scale factor for the image resolution used during feature detection and matching. Lower values (e.g., 0.6) are faster but may miss features in high-res images. 1.0 uses full resolution.") + seam_estimation_resol = gr.Slider(0.05, 1.0, step=0.05, value=0.1, label="Seam Estimation Resolution", + info="Scale factor for the image resolution used during seam finding (finding the optimal cut line). Lower values (e.g., 0.1) are much faster.") + compositing_resol = gr.Slider(-1.0, 1.0, step=0.1, value=-1.0, label="Compositing Resolution", + info="Scale factor for the image resolution used during the final blending stage. -1.0 uses the original source image resolution. Lower values reduce memory usage but might slightly blur the output.") + wave_correction = gr.Checkbox(value=False, label="Enable Wave Correction", + info="Attempts to correct perspective distortions (waviness) common in panoramas. Can increase processing time.") + exposure_comp_type = gr.Dropdown(["NO", "GAIN", "GAIN_BLOCKS"], value="GAIN_BLOCKS", label="Exposure Compensation", + info="Method used by the built-in stitcher to adjust brightness/contrast differences between images. 'GAIN_BLOCKS' is generally preferred for varying lighting.") + + # --- Detailed Stitcher Settings (Used for Video, DIRECT_PAIRWISE, and Fallback) --- + with gr.Accordion("Pairwise Stitching Settings (Video / Direct / Fallback)", open=True): + transform_model = gr.Radio(["Homography", "Affine_Partial", "Affine_Full"], label="Pairwise Transform Model", value="Homography", # Default to Homography + info="Geometric model for pairwise alignment. 'Homography' handles perspective. 'Affine' (Partial/Full) handles translation, rotation, scale, shear (better for scans, less distortion risk). If stitching fails with one model, try another.") + blend_method = gr.Radio(["Linear", "Multi-Band"], label="Blending Method", value="Multi-Band", + info="Algorithm for smoothing seams in overlapping regions when using the detailed stitcher (for video or image list fallback). 'Multi-Band' is often better but slower.") + enable_gain_compensation = gr.Checkbox(value=True, label="Enable Gain Compensation", + info="Adjusts overall brightness difference *before* blending when using the detailed stitcher. Recommended.") + orb_nfeatures = gr.Slider(500, 10000, step=100, value=2000, label="ORB Features", + info="Maximum ORB keypoints detected per image/frame. Used by the detailed stitcher (for video or image list fallback).") + match_ratio_thresh = gr.Slider(0.5, 0.95, step=0.01, value=0.75, label="Match Ratio Threshold", + info="Lowe's ratio test threshold for filtering feature matches (lower = stricter). Used by the detailed stitcher (for video or image list fallback).") + ransac_reproj_thresh = gr.Slider(1.0, 10.0, step=0.1, value=5.0, label="RANSAC Reproj Threshold", + info="Maximum reprojection error (pixels) allowed for a match to be considered an inlier by RANSAC during transformation estimation. Lower values are stricter.") + max_distance_coeff = gr.Slider(0.1, 2.0, step=0.05, value=0.5, label="Max Distance Coeff", + info="Multiplier for image diagonal used to filter initial matches. Limits the pixel distance between matched keypoints (0.5 means half the diagonal).") + max_blending_width = gr.Number(value=10000, label="Max Blending Width", precision=0, + info="Limits the canvas width during the detailed pairwise blending step to prevent excessive memory usage. Relevant for the detailed stitcher.") + max_blending_height = gr.Number(value=10000, label="Max Blending Height", precision=0, + info="Limits the canvas height during the detailed pairwise blending step to prevent excessive memory usage. Relevant for the detailed stitcher.") + blend_smooth_ksize = gr.Number(value=15, label="Blend Smooth Kernel Size", precision=0, + info="Size of Gaussian kernel to smooth blend mask/weights. Must be POSITIVE ODD integer to enable smoothing (e.g., 5, 15, 21). Set to -1 or an even number to disable smoothing.") + num_blend_levels = gr.Slider(2, 7, step=1, value=4, label="Multi-Band Blend Levels", + info="Number of pyramid levels for Multi-Band blending. Fewer levels are faster but might have less smooth transitions.") + + with gr.Accordion("Video Stitcher Settings", open=False): + sample_interval_ms = gr.Number(value=3000, label="Sample Interval (ms)", precision=0, + info="Time interval (in milliseconds) between sampled frames for video stitching. Smaller values sample more frames, increasing processing time but potentially improving tracking.") + max_composite_width_video = gr.Number(value=10000, label="Max Composite Width (Video)", precision=0, + info="Limits the width of the stitched output during video processing. If exceeded, the current result is saved and stitching restarts with the next frame. 0 = no limit.") + max_composite_height_video = gr.Number(value=10000, label="Max Composite Height (Video)", precision=0, + info="Limits the height of the stitched output during video processing. If exceeded, the current result is saved and stitching restarts with the next frame. 0 = no limit.") + + with gr.Accordion("Postprocessing Settings", open=False): + enable_cropping = gr.Checkbox(value=True, label="Crop Black Borders (Post-Stitch)", + info="Automatically remove black border areas from the final stitched image(s) AFTER stitching.") + strict_no_black_edges_checkbox = gr.Checkbox(value=False, label="Strict No Black Edges (Post-Crop)", + info="If 'Crop Black Borders' is enabled, this forces removal of *any* remaining black pixels directly on the image edges after the main crop. Might slightly shrink the image further.") + + with gr.Column(scale=1): + output_gallery = gr.Gallery( + label="Stitched Results", elem_id="output_gallery", object_fit="contain", type="filepath", rows=2, preview=True, height="auto", format="png", container=True) + output_log = gr.Textbox( + label="Status / Log", lines=20, interactive=False, show_copy_button=True) + + + # Define the list of inputs for the button click event + inputs=[ + input_files, + # Preprocessing + crop_top_percent, + crop_bottom_percent, + # OpenCV Stitcher (Image List) + stitcher_mode, # the selected string ("SCANS", "PANORAMA", or "DIRECT_PAIRWISE") + registration_resol, + seam_estimation_resol, + compositing_resol, + wave_correction, + exposure_comp_type, + # Postprocessing + enable_cropping, + strict_no_black_edges_checkbox, + # Detailed Stitcher Settings + transform_model, + blend_method, + enable_gain_compensation, + orb_nfeatures, + match_ratio_thresh, + ransac_reproj_thresh, + max_distance_coeff, + max_blending_width, + max_blending_height, + blend_smooth_ksize, + num_blend_levels, + # Video specific settings + sample_interval_ms, + max_composite_width_video, + max_composite_height_video + ] + + # Define examples (update to include the new transform_model parameter) + examples = [ + [ + ["examples/Wetter-Panorama/Wetter-Panorama1[NIuO6hrFTrg].mp4"], + 0, 20, + "DIRECT_PAIRWISE", 0.6, 0.1, -1, False, "GAIN_BLOCKS", + True, False, + "Homography", "Multi-Band", True, 5000, 0.5, 5.0, 0.5, 10000, 10000, 15, 4, + 2500, 10000, 10000, + ], + [ + ["examples/Wetter-Panorama/Wetter-Panorama2[NIuO6hrFTrg].mp4"], + 0, 20, + "DIRECT_PAIRWISE", 0.6, 0.1, -1, False, "GAIN_BLOCKS", + True, False, + "Homography", "Multi-Band", True, 5000, 0.5, 5.0, 0.5, 10000, 10000, 15, 4, + 2500, 10000, 10000, + ], + [ + ["examples/NieRAutomata/nier2B_01.jpg", "examples/NieRAutomata/nier2B_02.jpg", "examples/NieRAutomata/nier2B_03.jpg", "examples/NieRAutomata/nier2B_04.jpg", "examples/NieRAutomata/nier2B_05.jpg", + "examples/NieRAutomata/nier2B_06.jpg", "examples/NieRAutomata/nier2B_07.jpg", "examples/NieRAutomata/nier2B_08.jpg", "examples/NieRAutomata/nier2B_09.jpg", "examples/NieRAutomata/nier2B_10.jpg", ], + 0, 0, + "PANORAMA", 0.6, 0.1, -1, False, "GAIN_BLOCKS", + True, False, + "Homography", "Multi-Band", True, 5000, 0.5, 5.0, 0.5, 10000, 10000, 15, 4, + 5000, 10000, 10000, + ], + [ + ["examples/cat/cat_left.jpg", "examples/cat/cat_right.jpg"], + 0, 0, + "SCANS", 0.6, 0.1, -1, False, "GAIN_BLOCKS", + True, False, + "Affine_Partial", "Linear", True, 5000, 0.5, 5.0, 0.5, 10000, 10000, 15, 4, + 5000, 10000, 10000, + ], + [ + ["examples/ギルドの受付嬢ですが/Girumasu_1.jpg", "examples/ギルドの受付嬢ですが/Girumasu_2.jpg", "examples/ギルドの受付嬢ですが/Girumasu_3.jpg"], + 0, 0, + "PANORAMA", 0.6, 0.1, -1, False, "GAIN_BLOCKS", + True, False, + "Affine_Partial", "Linear", True, 5000, 0.65, 5.0, 0.5, 10000, 10000, 15, 4, + 5000, 10000, 10000, + ], + [ + ["examples/photographs1/img1.jpg", "examples/photographs1/img2.jpg", "examples/photographs1/img3.jpg", "examples/photographs1/img4.jpg"], + 0, 0, + "PANORAMA", 0.6, 0.1, -1, True, "GAIN_BLOCKS", + True, False, + "Homography", "Linear", True, 5000, 0.5, 5.0, 0.5, 10000, 10000, 15, 4, + 5000, 10000, 10000, + ] + ] + gr.Examples(examples, inputs=inputs, label="Example Configurations") + + # Connect button click to the function + stitch_button.click( + fn=run_stitching_interface, + inputs=inputs, + outputs=[output_gallery, output_log] + ) + +# --- Main Execution Block --- +if __name__ == "__main__": + print("Starting Gradio interface with selectable transformation model...") + # Enable queue for handling multiple requests and progress updates + demo.queue() + # Launch the interface + demo.launch(inbrowser=True) \ No newline at end of file