Coverage for core / src / sensorkit / webapi / preview.py: 92%

1# SPDX-License-Identifier: Apache-2.0 

2from __future__ import annotations 

3 

4import asyncio 

5import hashlib 

6import io 

7 

8_CACHE_MAX = 16 

9_cache: dict[bytes, asyncio.Task[PreviewJPEG]] = {} 

10 

11 

12def _discard_if_unsuccessful(checksum: bytes, task: asyncio.Task[PreviewJPEG]): 

13 """Drop failed/cancelled renders so a later request can retry instead of 

14 replaying the cached exception forever.""" 

15 if (task.cancelled() or task.exception() is not None) and _cache.get(checksum) is task: 

16 del _cache[checksum] 

17 

18 

19class PreviewJPEG: 

20 """A JPEG preview rendered from a FITS product, tagged with its source checksum. 

21 

22 Build instances with `from_fits`, which renders off the event loop and 

23 caches results by checksum so identical content is never re-rendered. Only the 

24 checksum and JPEG are retained -- never the raw FITS bytes. 

25 """ 

26 

27 def __init__(self, checksum: bytes, jpeg_bytes: bytes): 

28 self.checksum = checksum 

29 self.jpeg_bytes = jpeg_bytes 

30 

31 @classmethod 

32 async def from_fits(cls, raw_bytes: bytes) -> PreviewJPEG: 

33 """Return a preview for `raw_bytes`, rendering it off the event loop. 

34 

35 Cache bookkeeping runs on the event loop, so concurrent callers are 

36 serialized without a lock. Identical content in flight is rendered once 

37 and shared: callers await the same task. Only the hash and the 

38 astropy/PIL encode -- the CPU-bound work -- are offloaded to threads. 

39 """ 

40 checksum = (await asyncio.to_thread(hashlib.sha256, raw_bytes)).digest() 

41 task = _cache.get(checksum) 

42 

43 if task is not None: 

44 # Refresh recency: move to the end of the insertion order. 

45 del _cache[checksum] 

46 _cache[checksum] = task 

47 else: 

48 # A real Task (not a bare coroutine) so the shared render survives any 

49 # single caller being cancelled, e.g. a client disconnecting. 

50 task = asyncio.create_task(asyncio.to_thread(cls._encode, raw_bytes, checksum)) 

51 task.add_done_callback(lambda t: _discard_if_unsuccessful(checksum, t)) 

52 

53 # No await between insert and eviction, so the cache can't be observed 

54 # over-size or torn. The just-inserted task is newest, never evicted. 

55 _cache[checksum] = task 

56 

57 if len(_cache) > _CACHE_MAX: 

58 del _cache[next(iter(_cache))] 

59 

60 return await task 

61 

62 @classmethod 

63 def _encode(cls, raw: bytes, checksum: bytes): 

64 import numpy as np 

65 from astropy.io import fits 

66 from PIL import Image 

67 

68 with fits.open(io.BytesIO(raw)) as hdul: 

69 data = None 

70 

71 for hdu in hdul: 

72 if hdu.data is not None and hdu.data.ndim >= 2: 

73 data = hdu.data 

74 break 

75 

76 if data is None: 

77 raise ValueError("No 2D image data found in FITS file") 

78 

79 # Collapse any leading dimensions down to 2D 

80 while data.ndim > 2: 

81 data = data[0] 

82 

83 data = data.astype(np.float32) 

84 low, high = np.percentile(data, (1, 99)) 

85 

86 if high > low: 

87 scaled = np.clip((data - low) / (high - low), 0.0, 1.0) 

88 else: 

89 scaled = np.zeros_like(data) 

90 

91 img = Image.fromarray((scaled * 255).astype(np.uint8), mode="L") 

92 buf = io.BytesIO() 

93 img.save(buf, format="JPEG", quality=90) 

94 

95 return cls(checksum, buf.getvalue())