Skip to content

Instantly share code, notes, and snippets.

@grapeot
Created February 28, 2024 23:04
Show Gist options
  • Save grapeot/d3e92e5f382f962f0567af62d895933e to your computer and use it in GitHub Desktop.
Save grapeot/d3e92e5f382f962f0567af62d895933e to your computer and use it in GitHub Desktop.
A reference implementation of auto-focusing by local entropy. https://yage.ai/auto-focus.html
"""
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)
@siyu6974
Copy link

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

@grapeot
Copy link
Author

grapeot commented Feb 29, 2024

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

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

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment