Merge pull request #28 from mraspaud/feature-xarray-RGBA

Add support for RGBA from DataArrays

pyninjotiff/ninjotiff.py

@@ -438,14 +438,18 @@ def _finalize(img, dtype=np.uint8, value_range_measurement_unit=None,
 
  if img.mode == 'RGBA':
  if not isinstance(img, np.ma.MaskedArray):
- raise NotImplementedError("The 'RGBA' case has not been updated to xarray")
- channels, fill_value = img._finalize(dtype)
- fill_value = fill_value or (0, 0, 0, 0)
- data = np.dstack((channels[0].filled(fill_value[0]),
- channels[1].filled(fill_value[1]),
- channels[2].filled(fill_value[2]),
- channels[3].filled(fill_value[3])))
- return data, 1.0, 0.0, fill_value[0]
+ data, mode = img.finalize(fill_value=fill_value, dtype=dtype)
+ data = data.transpose('y', 'x', 'bands')
+ fill_value = fill_value or 0
+ else:
+ channels, fill_value = img._finalize(dtype)
+ fill_value = fill_value or (0, 0, 0, 0)
+ data = np.dstack((channels[0].filled(fill_value[0]),
+ channels[1].filled(fill_value[1]),
+ channels[2].filled(fill_value[2]),
+ channels[3].filled(fill_value[3])))
+ fill_value = fill_value[0]
+ return data, 1.0, 0.0, fill_value
 
  if img.mode == 'P':
  if not isinstance(img, np.ma.MaskedArray):

pyninjotiff/tests/test_ninjotiff.py

@@ -281,20 +281,8 @@ def test_write_rgb():
  """Test saving a non-trasparent RGB."""
  area = STEREOGRAPHIC_AREA
 
- x_size, y_size = 1024, 1024
- arr = np.zeros((3, y_size, x_size))
- radius = min(x_size, y_size) / 2.0
- centre = x_size / 2, y_size / 2
-
- for x in range(x_size):
- for y in range(y_size):
- rx = x - centre[0]
- ry = y - centre[1]
- s = ((x - centre[0])**2.0 + (y - centre[1])**2.0)**0.5 / radius
- if s <= 1.0:
- h = ((np.arctan2(ry, rx) / np.pi) + 1.0) / 2.0
- rgb = colorsys.hsv_to_rgb(h, s, 1.0)
- arr[:, y, x] = np.array(rgb)
+ fill_value = 0.0
+ arr = create_hsv_color_disk(fill_value)
 
  attrs = dict([('platform_name', 'NOAA-18'),
  ('resolution', 1050),
@@ -339,26 +327,30 @@ def test_write_rgb():
  arr[idx, :, :] * 255).astype(np.uint8))
 
 
-def test_write_rgb_with_a():
- """Test saving a transparent RGB."""
- area = STEREOGRAPHIC_AREA
-
+def create_hsv_color_disk(fill_value):
+ """Create an HSV colordisk."""
  x_size, y_size = 1024, 1024
- arr = np.zeros((3, y_size, x_size))
+ arr = np.full((3, y_size, x_size), fill_value)
  radius = min(x_size, y_size) / 2.0
  centre = x_size / 2, y_size / 2
-
  for x in range(x_size):
  for y in range(y_size):
  rx = x - centre[0]
  ry = y - centre[1]
- s = ((x - centre[0])**2.0 + (y - centre[1])**2.0)**0.5 / radius
+ s = ((x - centre[0]) ** 2.0 + (y - centre[1]) ** 2.0) ** 0.5 / radius
  if s <= 1.0:
  h = ((np.arctan2(ry, rx) / np.pi) + 1.0) / 2.0
  rgb = colorsys.hsv_to_rgb(h, s, 1.0)
  arr[:, y, x] = np.array(rgb)
- else:
- arr[:, y, x] = np.nan
+ return arr
+
+
+def test_write_rgb_with_a():
+ """Test saving a transparent RGB."""
+ area = STEREOGRAPHIC_AREA
+
+ fill_value = np.nan
+ arr = create_hsv_color_disk(fill_value)
 
  attrs = dict([('platform_name', 'NOAA-18'),
  ('resolution', 1050),
@@ -408,20 +400,8 @@ def test_write_rgb_tb():
  """Test saving a non-trasparent RGB with thumbnails."""
  area = STEREOGRAPHIC_AREA
 
- x_size, y_size = 1024, 1024
- arr = np.zeros((3, y_size, x_size))
- radius = min(x_size, y_size) / 2.0
- centre = x_size / 2, y_size / 2
-
- for x in range(x_size):
- for y in range(y_size):
- rx = x - centre[0]
- ry = y - centre[1]
- s = ((x - centre[0])**2.0 + (y - centre[1])**2.0)**0.5 / radius
- if s <= 1.0:
- h = ((np.arctan2(ry, rx) / np.pi) + 1.0) / 2.0
- rgb = colorsys.hsv_to_rgb(h, s, 1.0)
- arr[:, y, x] = np.array(rgb)
+ fill_value = 0.0
+ arr = create_hsv_color_disk(fill_value)
 
  attrs = dict([('platform_name', 'NOAA-18'),
  ('resolution', 1050),
@@ -592,6 +572,60 @@ def test_write_rgb_classified():
  np.testing.assert_allclose(res[:, :, 3] == 0, np.isnan(arr[0, :, :]))
 
 
+def test_write_rgba():
+ """Test saving an RGBA image."""
+ area = STEREOGRAPHIC_AREA
+
+ fill_value = np.nan
+ arr = create_hsv_color_disk(fill_value)
+
+ attrs = dict([('platform_name', 'NOAA-18'),
+ ('resolution', 1050),
+ ('polarization', None),
+ ('start_time', TIME - datetime.timedelta(minutes=55)),
+ ('end_time', TIME - datetime.timedelta(minutes=50)),
+ ('level', None),
+ ('sensor', 'avhrr-3'),
+ ('ancillary_variables', []),
+ ('area', area),
+ ('wavelength', None),
+ ('optional_datasets', []),
+ ('standard_name', 'overview'),
+ ('name', 'overview'),
+ ('prerequisites', [0.6, 0.8, 10.8]),
+ ('optional_prerequisites', []),
+ ('calibration', None),
+ ('modifiers', None),
+ ('mode', 'RGBA'),
+ ('enhancement_history', [{'scale': np.array([1, 1, -1]), 'offset': np.array([0, 0, 1])},
+ {'scale': np.array([0.0266347, 0.03559078, 0.01329783]),
+ 'offset': np.array([-0.02524969, -0.01996642, 3.8918446])},
+ {'gamma': 1.6}])])
+
+ kwargs = {'compute': True, 'fill_value': None, 'sat_id': 6300014,
+ 'chan_id': 6500015, 'data_cat': 'PPRN', 'data_source': 'SMHI', 'nbits': 8}
+ alpha = np.where(np.isnan(arr[0, :, :]), 0, 1)
+ arr = np.nan_to_num(arr)
+ arr = np.vstack((arr, alpha[np.newaxis, :, :]))
+ data = da.from_array(arr.clip(0, 1), chunks=1024)
+
+ data = xr.DataArray(data, coords={'bands': ['R', 'G', 'B', 'A']}, dims=[
+ 'bands', 'y', 'x'], attrs=attrs)
+
+ img = XRImage(data)
+ with tempfile.NamedTemporaryFile(delete=DELETE_FILES) as tmpfile:
+ filename = tmpfile.name
+ if not DELETE_FILES:
+ print(filename)
+ save(img, filename, data_is_scaled_01=True, **kwargs)
+ tif = TiffFile(filename)
+ res = tif[0].asarray()
+ for idx in range(4):
+ np.testing.assert_allclose(res[:, :, idx], np.round(
+ np.nan_to_num(arr[idx, :, :]) * 255).astype(np.uint8))
+ np.testing.assert_allclose(res[:, :, 3] == 0, alpha == 0)
+
+
 def test_write_bw_colormap():
  """Test saving a BW image with a colormap.