Coverage for core / src / sensorkit / data / focus.py: 0%
254 statements
« prev ^ index » next coverage.py v7.13.5, created at 2026-09-02 00:03 +0000
1# SPDX-License-Identifier: Apache-2.0
2import asyncio
3import io
4from typing import Literal
6import numpy as np
7from astropy.io import fits
8from pydantic import BaseModel, Field
10from sensorkit.common.keyword import declare_keyword
11from sensorkit.data.fits import ArrayInfo
12from sensorkit.data.graph import DataFlow, DataOp
15@declare_keyword
16class FocusInfo(BaseModel):
17 """Focus quality metrics for an image."""
18 hfr: float | None = Field(None, description="Half-Flux Radius in pixels")
19 fwhm: float | None = Field(None, description="Full Width Half Maximum in pixels")
20 star_count: int = Field(0, description="Number of detected stars")
21 focus_score: float | None = Field(None, description="Overall focus quality score (0-1, higher is better)")
24@declare_keyword
25class FocusFFTInfo(BaseModel):
26 """Focus quality metrics using Fourier Transform analysis."""
27 high_freq_ratio: float = Field(description="Ratio of high-frequency power to total power")
28 focus_score: float = Field(description="Focus quality score (0-1, higher is better)")
29 high_freq_threshold: float = Field(description="Threshold used to define high frequency region")
30 scale_factor: float = Field(description="Scale factor applied to normalize score")
33@declare_keyword
34class FocusConvolutionInfo(BaseModel):
35 """Focus quality metrics using convolution-based edge detection."""
36 method: str = Field(description="Edge detection method used (laplacian, sobel, variance_laplacian)")
37 raw_metric: float = Field(description="Raw focus metric value before normalization")
38 focus_score: float = Field(description="Focus quality score (0-1, higher is better)")
39 scale_factor: float = Field(description="Scale factor applied to normalize score")
42@declare_keyword
43class FocusConvolutionResults(BaseModel):
44 """Container for multiple convolution-based focus analysis results."""
45 results: dict[str, FocusConvolutionInfo] = Field(
46 default_factory=dict,
47 description="Focus analysis results keyed by method name"
48 )
50 def add_result(self, info: FocusConvolutionInfo):
51 """Add a focus analysis result."""
52 self.results[info.method] = info
55class AnalyzeFocusStars(DataOp):
56 """Analyze focus quality of an image and attach FocusInfo to context."""
57 op: Literal["analyze_focus_stars"] = "analyze_focus_stars"
58 detection_sigma: float = Field(5.0, description="Detection threshold in standard deviations above background")
59 min_star_flux: float = Field(100.0, description="Minimum integrated flux for star detection")
61 async def process(self, incoming: list[DataFlow], outgoing: list[DataFlow]):
62 context, buffer = await incoming[0].receive("buffer")
64 if buffer[:6] == b'SIMPLE':
65 image = await asyncio.to_thread(self._extract_fits_image, buffer)
66 else:
67 array_info = context.get(ArrayInfo)
68 if not array_info:
69 raise RuntimeError("analyze_focus requires ArrayInfo in context for raw array data")
70 image = await asyncio.to_thread(array_info.ndarray_from_buffer, buffer)
72 if image.ndim != 2:
73 raise RuntimeError("analyze_focus only supports 2D arrays")
75 focus_info = await asyncio.to_thread(self._analyze_focus, image)
77 context.set(focus_info)
79 await outgoing[0].send(context, buffer)
81 def _extract_fits_image(self, buffer: bytes) -> np.ndarray:
82 """Extract image data from a FITS buffer."""
83 with io.BytesIO(buffer) as bio:
84 with fits.open(bio) as hdul:
85 return hdul[0].data.astype(np.float64)
87 def _analyze_focus(self, image: np.ndarray) -> FocusInfo:
88 """Compute focus metrics from image data."""
89 try:
90 from scipy import ndimage
91 except ImportError:
92 return FocusInfo()
94 img = image.astype(np.float64)
96 background = np.median(img)
97 noise = np.std(img[img < np.percentile(img, 50)])
99 threshold = background + self.detection_sigma * noise
100 binary = img > threshold
101 labeled, num_features = ndimage.label(binary)
103 if num_features == 0:
104 return FocusInfo(star_count=0)
106 hfr_values = []
107 fwhm_values = []
108 star_count = 0
110 for label_id in range(1, num_features + 1):
111 star_metrics = self._analyze_star(img, labeled, label_id, background)
112 if star_metrics:
113 hfr_values.append(star_metrics['hfr'])
114 if star_metrics['fwhm'] is not None:
115 fwhm_values.append(star_metrics['fwhm'])
116 star_count += 1
118 return self._compute_focus_info(hfr_values, fwhm_values, star_count)
120 def _analyze_star(self, img: np.ndarray, labeled: np.ndarray, label_id: int, background: float) -> dict | None:
121 """Analyze a single star and return its metrics."""
122 mask = labeled == label_id
123 if not mask.any():
124 return None
126 positions = np.nonzero(mask)
127 fluxes = img[mask]
128 total_flux = fluxes.sum()
130 if total_flux < self.min_star_flux:
131 return None
133 centroid_y = np.average(positions[0], weights=fluxes)
134 centroid_x = np.average(positions[1], weights=fluxes)
136 hfr = self._calculate_hfr(positions, fluxes, centroid_y, centroid_x)
138 fwhm = self._calculate_fwhm(img, centroid_y, centroid_x, background)
140 return {'hfr': hfr, 'fwhm': fwhm}
142 def _calculate_hfr(self, positions: tuple, fluxes: np.ndarray, centroid_y: float, centroid_x: float) -> float:
143 """Calculate Half-Flux Radius for a star."""
144 distances = np.sqrt((positions[0] - centroid_y)**2 + (positions[1] - centroid_x)**2)
145 sorted_indices = np.argsort(distances)
146 sorted_fluxes = fluxes[sorted_indices]
147 cumulative_flux = np.cumsum(sorted_fluxes)
148 half_flux = fluxes.sum() / 2
150 hfr_idx = np.searchsorted(cumulative_flux, half_flux)
151 if hfr_idx < len(distances):
152 return float(distances[sorted_indices[hfr_idx]])
153 return 0.0
155 def _calculate_fwhm(self, img: np.ndarray, centroid_y: float, centroid_x: float, background: float) -> float | None:
156 """Calculate Full Width Half Maximum for a star."""
157 try:
158 cy, cx = int(centroid_y), int(centroid_x)
159 size = 15
160 y_min = max(0, cy - size)
161 y_max = min(img.shape[0], cy + size + 1)
162 x_min = max(0, cx - size)
163 x_max = min(img.shape[1], cx + size + 1)
165 star_region = img[y_min:y_max, x_min:x_max]
166 if star_region.size > 0:
167 peak = star_region.max()
168 half_max = (peak + background) / 2
169 half_max_pixels = star_region > half_max
170 if half_max_pixels.any():
171 area = half_max_pixels.sum()
172 return float(2 * np.sqrt(area / np.pi))
173 except (IndexError, ValueError):
174 pass
175 return None
177 def _compute_focus_info(self, hfr_values: list[float], fwhm_values: list[float], star_count: int) -> FocusInfo:
178 """Compute final focus metrics from collected star data."""
179 median_hfr = float(np.median(hfr_values)) if hfr_values else None
180 median_fwhm = float(np.median(fwhm_values)) if fwhm_values else None
182 focus_score = None
183 if median_hfr is not None:
184 focus_score = float(1.0 / (1.0 + median_hfr))
186 return FocusInfo(
187 hfr=median_hfr,
188 fwhm=median_fwhm,
189 star_count=star_count,
190 focus_score=focus_score
191 )
194class AnalyzeFocusFFT(DataOp):
195 """Analyze focus quality using Fourier Transform (frequency domain analysis).
197 Sharper images have more high-frequency content, so this method computes the
198 ratio of high-frequency power to total power in the image spectrum.
199 """
200 op: Literal["analyze_focus_fft"] = "analyze_focus_fft"
201 high_freq_threshold: float = Field(
202 0.3,
203 description="Fraction of frequency range considered 'high' (0-1)"
204 )
205 scale_factor: float = Field(
206 50.0,
207 description="Scale factor to normalize score to 0-1 range. Higher values for images with less high-freq content (astronomy: 50, satellites: 30)"
208 )
210 async def process(self, incoming: list[DataFlow], outgoing: list[DataFlow]):
211 context, buffer = await incoming[0].receive("buffer")
213 if buffer[:6] == b'SIMPLE':
214 image = await asyncio.to_thread(self._extract_fits_image, buffer)
215 else:
216 array_info = context.get(ArrayInfo)
217 if not array_info:
218 raise RuntimeError("analyze_focus_fft requires ArrayInfo in context for raw array data")
219 image = await asyncio.to_thread(array_info.ndarray_from_buffer, buffer)
221 if image.ndim != 2:
222 raise RuntimeError("analyze_focus_fft only supports 2D arrays")
224 focus_info = await asyncio.to_thread(self._analyze_focus_fft, image)
226 context.set(focus_info)
228 await outgoing[0].send(context, buffer)
230 def _extract_fits_image(self, buffer: bytes) -> np.ndarray:
231 """Extract image data from a FITS buffer."""
232 with io.BytesIO(buffer) as bio:
233 with fits.open(bio) as hdul:
234 return hdul[0].data.astype(np.float64)
236 def _analyze_focus_fft(self, image: np.ndarray) -> FocusFFTInfo:
237 """Compute focus metrics using FFT analysis."""
238 img = image.astype(np.float64)
240 # Normalize the image
241 img = (img - img.mean()) / (img.std() + 1e-8)
243 # Compute 2D FFT
244 fft = np.fft.fft2(img)
245 fft_shift = np.fft.fftshift(fft)
246 magnitude_spectrum = np.abs(fft_shift)
248 # Create frequency coordinates
249 rows, cols = img.shape
250 crow, ccol = rows // 2, cols // 2
252 # Create radial frequency map
253 y, x = np.ogrid[:rows, :cols]
254 r = np.sqrt((x - ccol)**2 + (y - crow)**2)
255 max_r = np.sqrt(crow**2 + ccol**2)
257 # Define high frequency region
258 high_freq_radius = self.high_freq_threshold * max_r
259 high_freq_mask = r > high_freq_radius
261 # Compute power in high frequencies vs total power
262 total_power = np.sum(magnitude_spectrum**2)
263 high_freq_power = np.sum((magnitude_spectrum * high_freq_mask)**2)
265 # Focus score based on high frequency content
266 # Normalized to 0-1 range using configurable scale factor
267 if total_power > 0:
268 high_freq_ratio = high_freq_power / total_power
269 focus_score = float(min(1.0, high_freq_ratio * self.scale_factor))
270 else:
271 high_freq_ratio = 0.0
272 focus_score = 0.0
274 return FocusFFTInfo(
275 high_freq_ratio=float(high_freq_ratio),
276 focus_score=focus_score,
277 high_freq_threshold=self.high_freq_threshold,
278 scale_factor=self.scale_factor
279 )
282class AnalyzeFocusConvolution(DataOp):
283 """Analyze focus quality using convolution-based edge detection.
285 Uses Laplacian or other edge detection kernels to measure sharpness.
286 Sharper images have stronger edge responses.
287 """
288 op: Literal["analyze_focus_convolution"] = "analyze_focus_convolution"
289 method: Literal["laplacian", "sobel", "variance_laplacian"] = Field(
290 "variance_laplacian",
291 description="Edge detection method to use"
292 )
293 scale_factor: float = Field(
294 1.0,
295 description="Divisor to normalize score to 0-1 range. Recommended: variance_laplacian=1.0, laplacian=0.5, sobel=1.0"
296 )
298 async def process(self, incoming: list[DataFlow], outgoing: list[DataFlow]):
299 context, buffer = await incoming[0].receive("buffer")
301 if buffer[:6] == b'SIMPLE':
302 image = await asyncio.to_thread(self._extract_fits_image, buffer)
303 else:
304 array_info = context.get(ArrayInfo)
305 if not array_info:
306 raise RuntimeError("analyze_focus_convolution requires ArrayInfo in context for raw array data")
307 image = await asyncio.to_thread(array_info.ndarray_from_buffer, buffer)
309 if image.ndim != 2:
310 raise RuntimeError("analyze_focus_convolution only supports 2D arrays")
312 focus_info = await asyncio.to_thread(self._analyze_focus_convolution, image)
314 # Get or create the results container
315 results = context.get(FocusConvolutionResults)
316 if results is None:
317 results = FocusConvolutionResults()
319 # Add this result to the container
320 results.add_result(focus_info)
321 context.set(results)
323 await outgoing[0].send(context, buffer)
325 def _extract_fits_image(self, buffer: bytes) -> np.ndarray:
326 """Extract image data from a FITS buffer."""
327 with io.BytesIO(buffer) as bio:
328 with fits.open(bio) as hdul:
329 return hdul[0].data.astype(np.float64)
331 def _analyze_focus_convolution(self, image: np.ndarray) -> FocusConvolutionInfo:
332 """Compute focus metrics using convolution-based edge detection."""
333 try:
334 from scipy import ndimage
335 except ImportError:
336 # Return a default result if scipy is not available
337 return FocusConvolutionInfo(
338 method=self.method,
339 raw_metric=0.0,
340 focus_score=0.0,
341 scale_factor=self.scale_factor
342 )
344 img = image.astype(np.float64)
346 # Normalize the image
347 img_norm = (img - img.mean()) / (img.std() + 1e-8)
349 if self.method == "laplacian":
350 # Laplacian kernel
351 laplacian = ndimage.laplace(img_norm)
352 focus_metric = np.abs(laplacian).mean()
354 elif self.method == "sobel":
355 # Sobel edge detection in both directions
356 sobel_x = ndimage.sobel(img_norm, axis=0)
357 sobel_y = ndimage.sobel(img_norm, axis=1)
358 edge_magnitude = np.sqrt(sobel_x**2 + sobel_y**2)
359 focus_metric = edge_magnitude.mean()
361 elif self.method == "variance_laplacian":
362 # Variance of Laplacian - excellent for focus detection
363 # This is essentially the Tenengrad operator
364 laplacian = ndimage.laplace(img_norm)
365 focus_metric = laplacian.var()
367 # Normalize focus metric to 0-1 range using configurable scale factor
368 focus_score = float(min(1.0, focus_metric / self.scale_factor))
370 return FocusConvolutionInfo(
371 method=self.method,
372 raw_metric=float(focus_metric),
373 focus_score=focus_score,
374 scale_factor=self.scale_factor
375 )
378class FocusInfoToFITS(DataOp):
379 """Add focus analysis keyword data to FITS headers.
381 Handles FocusInfo (star-based), FocusFFTInfo, and FocusConvolutionResults.
382 """
383 op: Literal["focus_info_to_fits"] = "focus_info_to_fits"
385 async def process(self, incoming: list[DataFlow], outgoing: list[DataFlow]):
386 context, buffer = await incoming[0].receive("buffer")
388 if buffer[:6] != b'SIMPLE':
389 raise RuntimeError("focus_info_to_fits requires FITS format buffer")
391 focus_info = context.get(FocusInfo)
392 fft_info = context.get(FocusFFTInfo)
393 conv_results = context.get(FocusConvolutionResults)
395 # If no focus data at all, pass through
396 if not any([focus_info, fft_info, conv_results]):
397 await outgoing[0].send(context, buffer)
398 return
400 updated_buffer = await asyncio.to_thread(
401 self._update_fits_headers, buffer, focus_info, fft_info, conv_results
402 )
403 await outgoing[0].send(context, updated_buffer)
405 @staticmethod
406 def _update_fits_headers(
407 buffer: bytes,
408 focus: FocusInfo | None,
409 fft: FocusFFTInfo | None,
410 conv: FocusConvolutionResults | None
411 ) -> bytes:
412 bio_in = io.BytesIO(buffer)
413 bio_out = io.BytesIO()
415 with fits.open(bio_in) as hdul:
416 header = hdul[0].header
418 # Star-based focus metrics
419 if focus:
420 if focus.hfr is not None:
421 header['FOCHFR'] = (focus.hfr, 'Half-Flux Radius in pixels')
422 if focus.fwhm is not None:
423 header['FOCFWHM'] = (focus.fwhm, 'Full Width Half Maximum in pixels')
424 if focus.star_count is not None:
425 header['FOCSTR'] = (focus.star_count, 'Number of detected stars')
426 if focus.focus_score is not None:
427 header['FOCSCR'] = (focus.focus_score, 'Focus score (star-based)')
429 # FFT-based focus metrics
430 if fft:
431 header['FOCFFT'] = (fft.focus_score, 'Focus score (FFT)')
432 header['FOCFFTR'] = (fft.high_freq_ratio, 'High-frequency ratio')
433 header['FOCFFTT'] = (fft.high_freq_threshold, 'FFT high-freq threshold')
435 # Convolution-based focus metrics
436 if conv and conv.results:
437 for method_name, info in conv.results.items():
438 # Use method abbreviations for FITS keywords (8 char limit)
439 prefix = {
440 'variance_laplacian': 'FOCVLAP',
441 'laplacian': 'FOCLAP',
442 'sobel': 'FOCSOB'
443 }.get(method_name, f'FOC{method_name[:3].upper()}')
445 header[prefix] = (info.focus_score, f'Focus score ({method_name})')
446 header[f'{prefix}R'] = (info.raw_metric, f'{method_name} raw metric')
448 hdul.writeto(bio_out)
450 return bio_out.getvalue()