stabilize.py

# I have 899 images in a folder, it is a picture of the sun against a black background.
# I would like to registrer these images to each other, so that they are all aligned.
# I will use OpenCV to do this.

# After registration I will crop the images so they are all the same size with the sun in the same place
# In a few images there is are some clouds, if so I would like to use the registration from the previous image.



import cv2
import os
import numpy as np
import matplotlib.pyplot as plt

def load_images_from_folder(folder):
    # Get all filenames
    all_filenames = os.listdir(folder)
    
    # Sort filenames properly using natural sorting
    # This ensures "1" comes before "2" and "2" comes before "10"
    def natural_sort_key(s):
        import re
        return [int(text) if text.isdigit() else text.lower() 
                for text in re.split(r'(\d+)', s)]
    
    filenames = sorted(all_filenames, key=natural_sort_key)
    
    images = []
    filename_map = []  # Keep track of the filenames for each image
    
    for filename in filenames:
        img_path = os.path.join(folder, filename)
        img = cv2.imread(img_path)
        if img is not None:
            images.append(img)
            filename_map.append(filename)
    
    return images, filename_map

def find_best_shift(reference_img, image, max_shift=10, shift_step=1, initial_shift=(0, 0)):
    """
    Find the best shift between two images using minimization of subtraction difference.
    
    Args:
        reference_img: Reference image
        image: Image to register
        max_shift: Maximum pixel shift to try in each direction
        shift_step: Step size for shift attempts
        initial_shift: Starting point for the shift search (used for manual refinement)
    
    Returns:
        best_shift: Tuple of (dx, dy) for the best shift
        best_score: Score of the best shift (lower is better)
        best_shifted_img: The shifted image with the best score
    """
    # Convert images to grayscale for registration
    gray_reference = cv2.cvtColor(reference_img, cv2.COLOR_BGR2GRAY)
    gray_image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    
    best_score = float('inf')  # Lower is better
    best_shift = initial_shift
    best_shifted_img = None
    
    # Calculate search range around the initial shift
    dx_start = initial_shift[0] - max_shift
    dx_end = initial_shift[0] + max_shift + 1
    dy_start = initial_shift[1] - max_shift
    dy_end = initial_shift[1] + max_shift + 1
    
    # Try different shifts to find the best match
    for dx in range(dx_start, dx_end, shift_step):
        for dy in range(dy_start, dy_end, shift_step):
            # Create transformation matrix
            M = np.float32([[1, 0, dx], [0, 1, dy]])
            
            # Apply shift
            shifted = cv2.warpAffine(gray_image, M, (gray_reference.shape[1], gray_reference.shape[0]))
            
            # Calculate difference between shifted image and reference
            diff = cv2.absdiff(gray_reference, shifted)
            score = np.sum(diff)  # Sum of absolute differences
            
            # Update best result if this shift is better
            if score < best_score:
                best_score = score
                best_shift = (dx, dy)
                # Store the shifted color image
                best_shifted_img = cv2.warpAffine(image, M, (image.shape[1], image.shape[0]))
    
    return best_shift, best_score, best_shifted_img


def register_images_by_subtraction(folder_path, output_folder=None, max_shift=10, shift_step=1, show_plots=True):
    """
    Registers images by finding the best shift that minimizes the subtraction difference,
    using the previous successfully registered image as reference.
    
    Args:
        folder_path: Path to folder containing images
        output_folder: Path to save registered images
        max_shift: Maximum pixel shift to try in each direction
        shift_step: Step size for shift attempts
        show_plots: Whether to show plots of each registration
    
    Returns:
        List of registered images and array of shifts
    """
    # Create output folder if specified
    if output_folder is not None:
        os.makedirs(output_folder, exist_ok=True)
    
    # Load images
    print(f"Loading images from {folder_path}...")
    images, filenames = load_images_from_folder(folder_path)
    
    if not images:
        print("No images found!")
        return [], []
    
    print(f"Loaded {len(images)} images")
    
    # Initialize lists to store results
    cumulative_shifts = [(0, 0)]  # Track cumulative shifts relative to first image
    registered_images = [images[0]]  # First image is the reference
    registration_quality = [True]  # Track if each image was well-registered
    
    # Save the first image if output folder is specified
    if output_folder is not None:
        output_path = os.path.join(output_folder, f"registered_0000.png")
        cv2.imwrite(output_path, images[0])

    # Keep track of the last well-registered image index
    last_good_ref_idx = 0
    
    # Process remaining images
    for i, (image, filename) in enumerate(zip(images[1:], filenames[1:]), 1):
        print(f"Processing image {i}/{len(images)-1}: {filename}")
        
        # Decide which reference image to use
        # If previous image was poorly registered, use the last known good one
        if i > 1 and not registration_quality[-1]:  # Make sure we only check after first image
            ref_idx = last_good_ref_idx
            print(f"  Using image {ref_idx} as reference (previous registration was poor)")
        else:
            ref_idx = i - 1  # Use previous image
        
        reference_img = images[ref_idx]  # Use original image as reference
        
        # Find the best shift for registration
        best_shift, best_score, temp_shifted_img = find_best_shift(
            reference_img, image, max_shift, shift_step
        )
        
        print(f"  Best shift: dx={best_shift[0]}, dy={best_shift[1]}, score={best_score}")
        
        # Check if the score is too high (indicating poor registration)
        score_threshold = 2500000  # Adjust based on your image characteristics
        is_good_registration = best_score <= score_threshold
        
        if not is_good_registration:
            print(f"  Warning: Poor registration quality, score={best_score}.")
            
            # Create a diff image for visualization
            gray_reference = cv2.cvtColor(reference_img, cv2.COLOR_BGR2GRAY)
            gray_registered = cv2.cvtColor(temp_shifted_img, cv2.COLOR_BGR2GRAY)
            diff_image = cv2.absdiff(gray_reference, gray_registered)
            
            # Create a figure for interactive selection
            fig = plt.figure(figsize=(15, 5))
            
            # Plot reference image
            ax1 = plt.subplot(1, 3, 1)
            plt.title(f"Reference (Image {ref_idx})")
            plt.imshow(cv2.cvtColor(reference_img, cv2.COLOR_BGR2RGB))
            plt.axis('off')
            
            # Plot registered image
            ax2 = plt.subplot(1, 3, 2)
            plt.title(f"Current (Image {i}) - Poor Registration")
            plt.imshow(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))  # Show original image for selection
            plt.axis('off')
            
            # Plot difference image
            plt.subplot(1, 3, 3)
            plt.title(f"Difference (Score: {best_score})")
            plt.imshow(diff_image, cmap='hot')
            plt.colorbar(label='Pixel Difference')
            plt.axis('off')
            
            plt.tight_layout()
            plt.suptitle(f"Registration Issue - Choose Action", fontsize=16)
            plt.subplots_adjust(top=0.85)
            
            # Add buttons for user choices - adjust positions for three buttons
            ax_use_prev = plt.axes([0.2, 0.01, 0.18, 0.05])
            ax_manual = plt.axes([0.4, 0.01, 0.18, 0.05])
            ax_use_best = plt.axes([0.6, 0.01, 0.18, 0.05])  # New button
            
            btn_use_prev = plt.Button(ax_use_prev, 'Use Previous Shift')
            btn_manual = plt.Button(ax_manual, 'Manual Points')
            btn_use_best = plt.Button(ax_use_best, 'Use Best Fit')  # New button
            
            # Store user choice
            user_choice = {'action': None, 'ref_point': None, 'img_point': None}
            
            def use_previous_shift(event):
                # Use the previous image's shift
                prev_shift = (0, 0)
                if i > 1:  # If not the first image after reference
                    # Calculate shift between previous images
                    prev_shift = (
                        cumulative_shifts[i-1][0] - cumulative_shifts[i-2][0],
                        cumulative_shifts[i-1][1] - cumulative_shifts[i-2][1]
                    )
                user_choice['action'] = 'prev_shift'
                user_choice['shift'] = prev_shift
                plt.close(fig)
            
            def enable_manual_selection(event):
                user_choice['action'] = 'manual'
                # Change cursor to indicate selection mode
                plt.suptitle(f"Click on a feature in reference image (left), then same feature in current image (right)", fontsize=14)
                plt.draw()
            
            def use_best_fit(event):
                # Use the calculated best shift despite poor quality
                user_choice['action'] = 'use_best'
                plt.close(fig)
            
            # Register clicks for manual point selection
            def onclick(event):
                if user_choice['action'] != 'manual':
                    return
                
                if event.inaxes == ax1:  # Reference image
                    user_choice['ref_point'] = (event.xdata, event.ydata)
                    plt.suptitle(f"Reference point selected! Now click the same point in current image", fontsize=14)
                    plt.draw()
                elif event.inaxes == ax2 and user_choice['ref_point'] is not None:  # Current image and ref point selected
                    user_choice['img_point'] = (event.xdata, event.ydata)
                    plt.suptitle(f"Both points selected! Processing...", fontsize=14)
                    plt.draw()
                    plt.pause(0.5)  # Brief pause to show message
                    plt.close(fig)
            
            # Connect event handlers
            btn_use_prev.on_clicked(use_previous_shift)
            btn_manual.on_clicked(enable_manual_selection)
            btn_use_best.on_clicked(use_best_fit)  # Connect new button
            fig.canvas.mpl_connect('button_press_event', onclick)
            
            plt.show()
            
            # Process based on user choice
            if user_choice['action'] == 'prev_shift':
                # Use the previous shift
                best_shift = user_choice['shift']
                print(f"  Using previous shift: dx={best_shift[0]}, dy={best_shift[1]}")
                
                # Create shifted image with this shift
                M = np.float32([[1, 0, best_shift[0]], [0, 1, best_shift[1]]])
                temp_shifted_img = cv2.warpAffine(image, M, (image.shape[1], image.shape[0]))
                
            elif user_choice['action'] == 'manual' and user_choice['ref_point'] is not None and user_choice['img_point'] is not None:
                # Calculate manual shift
                ref_pt = user_choice['ref_point']
                img_pt = user_choice['img_point']
                manual_dx = int(ref_pt[0] - img_pt[0])
                manual_dy = int(ref_pt[1] - img_pt[1])
                
                print(f"  Manual shift starting point: dx={manual_dx}, dy={manual_dy}")
                
                # Rerun registration with the manual points as starting position
                refined_shift, refined_score, temp_shifted_img = find_best_shift(
                    reference_img, image, max_shift=5, shift_step=1, 
                    initial_shift=(manual_dx, manual_dy)
                )
                
                best_shift = refined_shift
                best_score = refined_score
                print(f"  Refined shift: dx={best_shift[0]}, dy={best_shift[1]}, score={best_score}")
                
                # Check if manual registration is good
                is_good_registration = True  # Assume manual selection is good
            
            elif user_choice['action'] == 'use_best':
                # Keep the calculated best shift but mark as poor registration
                print(f"  Using calculated best shift despite poor quality: dx={best_shift[0]}, dy={best_shift[1]}")
                # is_good_registration remains False - we don't update reference image
            
            else:
                # If no valid choice made, keep original result
                print("  No valid action selected. Using original poor registration.")
        
        else:
            # Registration is good, update the last good reference index
            last_good_ref_idx = i
            
            # Generate plots if needed
            if show_plots and output_folder is not None:
                # Calculate the difference between reference and registered image for display
                gray_reference = cv2.cvtColor(reference_img, cv2.COLOR_BGR2GRAY)
                gray_registered = cv2.cvtColor(temp_shifted_img, cv2.COLOR_BGR2GRAY)
                diff_image = cv2.absdiff(gray_reference, gray_registered)
                
                # Create a figure with 1 row and 3 columns
                fig = plt.figure(figsize=(15, 5))
                
                # Plot reference image
                plt.subplot(1, 3, 1)
                plt.title(f"Reference (Image {ref_idx})")
                plt.imshow(cv2.cvtColor(reference_img, cv2.COLOR_BGR2RGB))
                plt.axis('off')
                
                # Plot registered image
                plt.subplot(1, 3, 2)
                plt.title(f"Registered (Image {i}) - Good")
                plt.imshow(cv2.cvtColor(temp_shifted_img, cv2.COLOR_BGR2RGB))
                plt.axis('off')
                
                # Plot difference image
                plt.subplot(1, 3, 3)
                plt.title(f"Difference (Score: {best_score})")
                plt.imshow(diff_image, cmap='hot')
                plt.colorbar(label='Pixel Difference')
                plt.axis('off')
                
                plt.tight_layout()
                plt.suptitle(f"Registration Results for Image {i}", fontsize=16)
                plt.subplots_adjust(top=0.85)
                
                # Save the plot if output folder is specified
                plot_dir = os.path.join(output_folder, "plots")
                os.makedirs(plot_dir, exist_ok=True)
                plt.savefig(os.path.join(plot_dir, f"registration_plot_{i:04d}.png"))
                
                # Close the figure without showing it
                plt.close(fig)
        
        # Calculate cumulative shift relative to first image
        # If we used a non-previous image as reference, adjust the calculation
        if ref_idx != i - 1:
            # Calculate shift relative to the reference image we used
            ref_cumulative_shift = cumulative_shifts[ref_idx]
            current_cumulative_shift = (
                ref_cumulative_shift[0] + best_shift[0],
                ref_cumulative_shift[1] + best_shift[1]
            )
        else:
            # Normal case: shift relative to previous image
            prev_cumulative_shift = cumulative_shifts[-1]
            current_cumulative_shift = (
                prev_cumulative_shift[0] + best_shift[0],
                prev_cumulative_shift[1] + best_shift[1]
            )
        
        # Apply the cumulative shift to the original image
        M = np.float32([[1, 0, current_cumulative_shift[0]], [0, 1, current_cumulative_shift[1]]])
        registered_img = cv2.warpAffine(image, M, (image.shape[1], image.shape[0]))
        
        # Store the results
        registered_images.append(registered_img)
        cumulative_shifts.append(current_cumulative_shift)
        registration_quality.append(is_good_registration)
        
        # Save the registered image if requested
        if output_folder is not None:
            output_path = os.path.join(output_folder, f"registered_{i:04d}.png")
            cv2.imwrite(output_path, registered_img)
    
    print("Registration complete!")
    return registered_images, cumulative_shifts

def main():
    # Update with your folder paths
    input_folder = r"Y:\Timelapse\2025-03-29 - Solformorkelse\Solformorkelse - TIFF"
    output_folder = r"Y:\Timelapse\2025-03-29 - Solformorkelse\Solformorkelse - Koregistrert"
    
    # Register and crop images
    registered_images, shifts = register_images_by_subtraction(
        input_folder, 
        output_folder,
        max_shift=20,
        shift_step=1,
        show_plots=True  # Set to False if you don't want plots for all 899 images
    )
    
    # Print the shift information
    print("\nShift information:")
    for i, (dx, dy) in enumerate(shifts):
        print(f"Image {i}: dx={dx}, dy={dy}")
    
    # Display first and last registered image
    if len(registered_images) >= 2:
        plt.figure(figsize=(12, 6))
        
        plt.subplot(1, 2, 1)
        plt.title("First Image")
        plt.imshow(cv2.cvtColor(registered_images[0], cv2.COLOR_BGR2RGB))
        
        plt.subplot(1, 2, 2)
        plt.title("Last Registered Image")
        plt.imshow(cv2.cvtColor(registered_images[-1], cv2.COLOR_BGR2RGB))
        
        plt.tight_layout()
        plt.show()

if __name__ == "__main__":
    main()