A reference implementation of auto-focusing by local entropy. https://yage.ai/auto-focus.html

auto_focus.py

"""
A class supporting the auto-focusing of a telescope.
"""
from astropy.io import fits
from astropy.stats import sigma_clipped_stats
from photutils.detection import DAOStarFinder
from glob import glob
from tqdm import tqdm
from typing import List, Dict, Any
from time import time
import numpy as np
import numpy.typing as npt
import matplotlib.pyplot as plt
import logging
import re
from random import shuffle
from multiprocessing import Pool, cpu_count
    
def extract_sources_from_image_data(img_data: npt.ArrayLike, step: int) -> Dict[str, Any]:
    """
    Extracts the sources from some image data.

    Args:
        - image_data: the image data to extract the sources from.
        - step: the step of the focuser.
    
    Returns:
        - a dictionary containing the step and the sources.
    """
    mean, median, std = sigma_clipped_stats(img_data, sigma=3.0)  
    finder = DAOStarFinder(fwhm=3.0, threshold=max(90, 3*std))  
    sources = finder(img_data - median)
    return {'step': step, 'sources': sources}

def auto_focus(fns: List[str]) -> int:
    """
    A testing utility function.
    """
    def read_fit_file(fn: str) -> npt.ArrayLike:
        return fits.open(fn)[0].data
    def parse_step(fn: str) -> int:
        return int(re.search(r'auto_focus_(\d+)_3s.fits', fn).group(1))
    img_data = [read_fit_file(fn) for fn in tqdm(fns)]
    steps = [parse_step(fn) for fn in fns]
    return calculate_focus_step(img_data, steps)

def calculate_focus_step(image_data: List[npt.ArrayLike], steps: List[int]) -> int:
    """
    Calculates the focus step from the image data and the focus steps.
    This function assumes the image_data is from the raw image, and will crop it to the central area.
    The algorithm works by first extracting the sources from the image data, then calculating the entropy of the local region around the sources.

    Args:
        - image_data: the image data to calculate the focus step from.
        - focus_steps: the focus steps corresponding to the image data.

    Returns:
        - the focus step that is likely the focus point.
    """
    pool = Pool(min(cpu_count(), 24))
    # Only keep the central area for auto focus
    # Then extract sources from each image
    central_regions = []
    for i, img_data in enumerate(image_data):
        h, w = img_data.shape[:2]
        central_region = img_data[h//4:3*h//4, w//4:3*w//4]
        central_regions.append(central_region)
    data = pool.starmap(extract_sources_from_image_data, zip(central_regions, steps))
    data = sorted(data, key=lambda x: x['step'])
    data = [d['sources'].to_pandas() for d in data]

    # Use the union of the star locations to calculate the entropy of the local region around the stars
    logging.info('Collecting stars...')
    star_locs = []
    for df in data:
        for i, r in df.iterrows():
            star_locs.append((i, r['xcentroid'], r['ycentroid'], 0))
    logging.info(f'Found {len(star_locs)} stars.')
    # Cap the number of stars to 5000
    if len(star_locs) > 5000:
        shuffle(star_locs)
        star_locs = star_locs[:5000]

    # For each image, we want to calculate the average entropy of the local regions around the stars
    local_region_size = 30
    image_entropies = []
    for i, d in enumerate(data):
        image = central_regions[i]
        entropies = []
        for _, x, y, peak in star_locs:
            x = int(x)
            y = int(y)
            local_region = image[y - local_region_size:y + local_region_size, x - local_region_size:x + local_region_size]
            # Calculate the entropy of the local region, by first calculating a histogram up to brightness 1000
            hist, _ = np.histogram(local_region, bins=1000)
            if np.sum(hist) == 0:
                continue
            hist = hist / np.sum(hist)
            entropies.append(-np.sum(hist * np.log(hist + 1e-10)))
        image_entropies.append(np.mean(entropies))
    peak_index = np.argmin(image_entropies)

    # Fit a 3rd degree polynomial to the data, with weights as a Gaussian of the entropy from the min entropy
    min_step = np.min(steps)
    max_step = np.max(steps)
    x = np.linspace(min_step, max_step, 1000)
    min_entropy = np.min(image_entropies)
    # Calculate the weights
    weights = np.exp(-(image_entropies - min_entropy)**2 / (3 * np.std(image_entropies)**2))
    fitted_curve = np.poly1d(np.polyfit(steps, image_entropies, 3, w=weights))
    y3 = fitted_curve(x)
    min_x_range = np.linspace(steps[peak_index] - 100, steps[peak_index] + 100, 100)
    min_y_range = fitted_curve(min_x_range)
    min_x = np.argmin(min_y_range)
    min_x = min_x_range[min_x]
    min_y = fitted_curve(min_x)
    plt.figure()
    plt.plot(steps, image_entropies, 'o', linestyle='--', label='Data')
    plt.plot(x, y3, label='Weighted')
    plt.plot(min_x, min_y, 'o', label='Minimum')
    plt.legend()
    plt.savefig(f'auto_focus_{time()}.png')
    plt.close()

    return int(min_x)

if __name__ == '__main__':
    logging.basicConfig(format='%(asctime)s : %(levelname)s : %(message)s', level=logging.INFO)
    # calculate_focus_step([], [])
    fns = glob('auto_focus_*.fits')
    fns = [fn for fn in fns if re.search(r'auto_focus_1\d{4}_3s.fits', fn)]
    auto_focus(fns)

鸭哥可以考虑用 sep 替换 DAOStarFinder，我的镜子上测试后感觉 sep 更加robust

鸭哥可以考虑用 sep 替换 DAOStarFinder，我的镜子上测试后感觉 sep 更加robust

谢谢！我试了一下也调了一些参数，感觉有些场景效果不是很好，可能是我调参还没调好。回头有时间具体看看。