Added extra type hints

This commit is contained in:
Joel Collins 2020-12-04 15:23:03 +00:00
parent 1de6b2a0ab
commit 7ba3f447e1

View file

@ -8,14 +8,16 @@ from picamerax import PiCamera
from picamerax.array import PiBayerArray, PiRGBArray from picamerax.array import PiBayerArray, PiRGBArray
def rgb_image(camera: PiCamera, resize: Optional[Tuple[int, int]] = None, **kwargs): def rgb_image(
camera: PiCamera, resize: Optional[Tuple[int, int]] = None, **kwargs
) -> PiRGBArray:
"""Capture an image and return an RGB numpy array""" """Capture an image and return an RGB numpy array"""
with PiRGBArray(camera, size=resize) as output: with PiRGBArray(camera, size=resize) as output:
camera.capture(output, format="rgb", resize=resize, **kwargs) camera.capture(output, format="rgb", resize=resize, **kwargs)
return output.array return output.array
def flat_lens_shading_table(camera: PiCamera): def flat_lens_shading_table(camera: PiCamera) -> np.ndarray:
"""Return a flat (i.e. unity gain) lens shading table. """Return a flat (i.e. unity gain) lens shading table.
This is mostly useful because it makes it easy to get the size This is mostly useful because it makes it easy to get the size
@ -107,9 +109,13 @@ def lst_from_channels(channels: np.ndarray) -> np.ndarray:
[channels.shape[0]] + lst_resolution, dtype=np.float [channels.shape[0]] + lst_resolution, dtype=np.float
) )
for i in range(lens_shading.shape[0]): for i in range(lens_shading.shape[0]):
image_channel = channels[i, :, :] image_channel: np.ndarray = channels[i, :, :]
iw: int
ih: int
iw, ih = image_channel.shape iw, ih = image_channel.shape
ls_channel = lens_shading[i, :, :] ls_channel: np.ndarray = lens_shading[i, :, :]
lw: int
lh: int
lw, lh = ls_channel.shape lw, lh = ls_channel.shape
# The lens shading table is rounded **up** in size to 1/64th of the size of # The lens shading table is rounded **up** in size to 1/64th of the size of
# the image. Rather than handle edge images separately, I'm just going to # the image. Rather than handle edge images separately, I'm just going to
@ -118,7 +124,7 @@ def lst_from_channels(channels: np.ndarray) -> np.ndarray:
# half the size of the full image - remember the Bayer pattern... This # half the size of the full image - remember the Bayer pattern... This
# should give results very close to 6by9's solution, albeit considerably # should give results very close to 6by9's solution, albeit considerably
# less computationally efficient! # less computationally efficient!
padded_image_channel = np.pad( padded_image_channel: np.ndarray = np.pad(
image_channel, [(0, lw * 32 - iw), (0, lh * 32 - ih)], mode="edge" image_channel, [(0, lw * 32 - iw), (0, lh * 32 - ih)], mode="edge"
) # Pad image to the right and bottom ) # Pad image to the right and bottom
logging.info( logging.info(
@ -131,7 +137,7 @@ def lst_from_channels(channels: np.ndarray) -> np.ndarray:
) )
# Next, fill the shading table (except edge pixels). Please excuse the # Next, fill the shading table (except edge pixels). Please excuse the
# for loop - I know it's not fast but this code needn't be! # for loop - I know it's not fast but this code needn't be!
box = 3 # We average together a square of this side length for each pixel. box: int = 3 # We average together a square of this side length for each pixel.
# NB this isn't quite what 6by9's program does - it averages 3 pixels # NB this isn't quite what 6by9's program does - it averages 3 pixels
# horizontally, but not vertically. # horizontally, but not vertically.
for dx in np.arange(box) - box // 2: for dx in np.arange(box) - box // 2:
@ -152,10 +158,10 @@ def lst_from_channels(channels: np.ndarray) -> np.ndarray:
# What we actually want to calculate is the gains needed to compensate for the # What we actually want to calculate is the gains needed to compensate for the
# lens shading - that's 1/lens_shading_table_float as we currently have it. # lens shading - that's 1/lens_shading_table_float as we currently have it.
gains = 32.0 / lens_shading # 32 is unity gain gains: np.ndarray = 32.0 / lens_shading # 32 is unity gain
gains[gains > 255] = 255 # clip at 255, maximum gain is 255/32 gains[gains > 255] = 255 # clip at 255, maximum gain is 255/32
gains[gains < 32] = 32 # clip at 32, minimum gain is 1 (is this necessary?) gains[gains < 32] = 32 # clip at 32, minimum gain is 1 (is this necessary?)
lens_shading_table = gains.astype(np.uint8) lens_shading_table: np.ndarray = gains.astype(np.uint8)
return lens_shading_table[::-1, :, :].copy() return lens_shading_table[::-1, :, :].copy()