Smarter about choosing right color mode/bit depth

lib/ex_png/chunks/image_data.ex

@@ -8,6 +8,11 @@ defmodule ExPng.Chunks.ImageData do
   decoded by `ExPng`.
   """
 
+  use ExPng.Constants
+
+  alias ExPng.Image
+  import ExPng.Utilities, only: [reduce_to_binary: 1]
+
   @type t :: %__MODULE__{
     data: binary,
     type: :IDAT
@@ -49,31 +54,64 @@ defmodule ExPng.Chunks.ImageData do
     <<length::32>> <> type <> data <> <<crc::32>>
   end
 
-  def from_pixels(pixels) do
+  def from_pixels(image, header) do
+    palette = Image.unique_pixels(image)
     data =
-      Enum.map(pixels, fn line ->
-        Task.async(fn -> line_to_binary(line) end)
+      Enum.map(image.pixels, fn line ->
+        Task.async(fn -> line_to_binary(line, header, palette) end)
       end)
       |> Enum.map(fn task ->
         Task.await(task)
       end)
-      |> Enum.reverse()
-      |> Enum.reduce(&Kernel.<>/2)
+      |> reduce_to_binary()
 
-    %__MODULE__{data: data}
+    {%__MODULE__{data: data}, %ExPng.Chunks.Palette{palette: palette}}
   end
 
   ## PRIVATE
 
-  defp line_to_binary(line) do
-    Enum.reduce(line, <<0>>, fn pixel, acc ->
-      acc <> <<pixel.r, pixel.g, pixel.b, pixel.a>>
-    end)
+  defp line_to_binary(line, %{color_mode: @indexed} = header, palette) do
+    bit_depth = header.bit_depth
+    chunk_size = div(8, bit_depth)
+    line =
+      line
+      |> Enum.map(fn pixel -> Enum.find_index(palette, fn p -> p == pixel end) end)
+      |> Enum.map(fn i ->
+        Integer.to_string(i, 2)
+        |> String.pad_leading(bit_depth, "0")
+      end)
+      |> Enum.chunk_every(chunk_size, chunk_size)
+      |> Enum.map(fn byte ->
+        byte =
+          byte
+          |> Enum.join("")
+          |> String.pad_trailing(8, "0")
+          |> String.to_integer(2)
+        <<byte>>
+      end)
+      |> reduce_to_binary()
+    <<0>> <> line
+  end
+
+  defp line_to_binary(line, %{bit_depth: 1} = _header, _palette) do
+    line =
+      line
+      |> Enum.map(& div(&1.b, 255))
+      |> Enum.chunk_every(8)
+      |> Enum.map(fn bits -> Enum.join(bits, "") |> String.to_integer(2) end)
+      |> Enum.map(fn byte -> <<byte>> end)
+      |> reduce_to_binary()
+    <<0>> <> line
   end
 
-  defp reduce_to_binary(chunks) do
-    Enum.reduce(chunks, <<>>, fn chunk, acc ->
-      acc <> chunk
+  defp line_to_binary(line, %{bit_depth: 8} = header, _palette) do
+    Enum.reduce(line, <<0>>, fn pixel, acc ->
+      acc <> case header.color_mode do
+        @truecolor_alpha -> <<pixel.r, pixel.g, pixel.b, pixel.a>>
+        @truecolor -> <<pixel.r, pixel.g, pixel.b>>
+        @grayscale_alpha -> <<pixel.b, pixel.a>>
+        @grayscale -> <<pixel.b>>
+      end
     end)
   end
 
@@ -94,7 +132,6 @@ defmodule ExPng.Chunks.ImageData do
     :zlib.close(zstream)
     deflated_data
     |> List.flatten()
-    |> Enum.reverse()
-    |> Enum.reduce(&Kernel.<>/2)
+    |> reduce_to_binary()
   end
 end

lib/ex_png/chunks/palette.ex

@@ -6,6 +6,8 @@ defmodule ExPng.Chunks.Palette do
 
   alias ExPng.Pixel
 
+  import ExPng.Utilities, only: [reduce_to_binary: 1]
+
   @type t :: %__MODULE__{
     type: :PLTE,
     data: binary(),
@@ -20,6 +22,20 @@ defmodule ExPng.Chunks.Palette do
     end
   end
 
+  @behaviour ExPng.Encodeable
+
+  @impl true
+  def to_bytes(%__MODULE__{palette: palette}, _encoding_options \\ []) do
+    data =
+      Enum.map(palette, fn pixel -> <<pixel.r, pixel.g, pixel.b>> end)
+      |> reduce_to_binary()
+    length = byte_size(data)
+    type = <<80, 76, 84, 69>>
+    crc = :erlang.crc32([type, data])
+
+    <<length::32>> <> type <> data <> <<crc::32>>
+  end
+
   ## PRIVATE
 
   defp parse_palette(data) do

lib/ex_png/image.ex

@@ -77,6 +77,13 @@ defmodule ExPng.Image do
     end
   end
 
+  @spec unique_pixels(__MODULE__.t) :: integer()
+  def unique_pixels(%__MODULE__{pixels: pixels}) do
+    pixels
+    |> List.flatten()
+    |> Enum.uniq()
+  end
+
   defdelegate erase(image), to: Drawing
   defdelegate draw(image, coordinates, color), to: Drawing
   defdelegate at(image, coordinates), to: Drawing

lib/ex_png/image/decoding.ex

@@ -34,7 +34,7 @@ defmodule ExPng.Image.Decoding do
           image.header_chunk.bit_depth,
           image.header_chunk.color_mode,
           image.palette_chunk
-        )
+        ) |> Enum.take(data.header_chunk.width)
       end)
 
     %Image{

lib/ex_png/image/encoding.ex

@@ -7,31 +7,97 @@ defmodule ExPng.Image.Encoding do
   use ExPng.Constants
 
   alias ExPng.Chunks.{End, Header, ImageData}
-  alias ExPng.{Image, RawData}
+  alias ExPng.{Image, Pixel, RawData}
 
   def to_raw_data(%Image{} = image) do
     header = build_header(image)
-    image_data_chunk = ImageData.from_pixels(image.pixels)
+    {image_data_chunk, palette} = ImageData.from_pixels(image, header)
 
-    {
-      :ok,
+    raw_data =
       %RawData{
         header_chunk: header,
         data_chunk: image_data_chunk,
         end_chunk: %End{}
       }
-    }
+
+    raw_data = case header.color_mode do
+      @indexed ->
+        %{raw_data |
+          palette_chunk: palette
+        }
+      _ -> raw_data
+    end
+
+    {:ok, raw_data}
   end
 
   defp build_header(%Image{} = image) do
+    {bit_depth, color_mode} = bit_depth_and_color_mode(image)
     %Header{
       width: image.width,
       height: image.height,
-      bit_depth: 8,
-      color_mode: @truecolor_alpha,
+      bit_depth: bit_depth,
+      color_mode: color_mode,
       compression: 0,
       filter: 0,
       interlace: 0
     }
   end
+
+  # 1. black and white
+  # 2. grayscale (alpha)
+  # 3. indexed
+  # 4. truecolor (alpha)
+  defp bit_depth_and_color_mode(%Image{} = image) do
+    case black_and_white?(image) do
+      true -> {1, @grayscale}
+      false ->
+        case {indexable?(image), opaque?(image), grayscale?(image)} do
+          {_, true, true} -> {8, @grayscale}
+          {_, false, true} -> {8, @grayscale_alpha}
+          {true, _, _} -> {indexed_bit_depth(image), @indexed}
+          {_, true, false} -> {8, @truecolor}
+          {_, false, false} -> {8, @truecolor_alpha}
+        end
+    end
+  end
+
+  defp indexed_bit_depth(%Image{} = image) do
+    case unique_pixel_count(image) do
+      i when i <= 2 -> 1
+      i when i <= 4 -> 2
+      i when i <= 16 -> 4
+      i when i <= 256 -> 8
+    end
+  end
+
+  defp indexable?(%Image{} = image) do
+    image
+    |> unique_pixel_count()
+    |> Kernel.<=(256)
+  end
+
+  defp unique_pixel_count(%Image{} = image) do
+    image
+    |> Image.unique_pixels()
+    |> length()
+  end
+
+  defp opaque?(%Image{pixels: pixels}) do
+    pixels
+    |> List.flatten()
+    |> Enum.all?(&Pixel.opaque?/1)
+  end
+
+  defp grayscale?(%Image{pixels: pixels}) do
+    pixels
+    |> List.flatten()
+    |> Enum.all?(&Pixel.grayscale?/1)
+  end
+
+  defp black_and_white?(%Image{pixels: pixels}) do
+    pixels
+    |> List.flatten()
+    |> Enum.all?(&Pixel.black_or_white?/1)
+  end
 end

lib/ex_png/pixel.ex

@@ -61,6 +61,19 @@ defmodule ExPng.Pixel do
   """
   @spec white() :: __MODULE__.t()
   def white, do: grayscale(255)
+
+  @spec opaque?(__MODULE__.t()) :: boolean
+  def opaque?(%__MODULE__{a: 255}), do: true
+  def opaque?(_), do: false
+
+  @spec grayscale?(__MODULE__.t()) :: boolean
+  def grayscale?(%__MODULE__{r: gr, g: gr, b: gr}), do: true
+  def grayscale?(_), do: false
+
+  @spec black_or_white?(__MODULE__.t()) :: boolean
+  def black_or_white?(%__MODULE__{r: 0, g: 0, b: 0, a: 255}), do: true
+  def black_or_white?(%__MODULE__{r: 255, g: 255, b: 255, a: 255}), do: true
+  def black_or_white?(_), do: false
 end
 
 defimpl Inspect, for: ExPng.Pixel do

lib/ex_png/raw_data.ex

@@ -52,9 +52,15 @@ defmodule ExPng.RawData do
   def to_file(%__MODULE__{} = raw_data, filename, encoding_options \\ []) do
     image_data = ImageData.to_bytes(raw_data.data_chunk, encoding_options)
 
+    palette_data = case raw_data.palette_chunk do
+      nil -> ""
+      palette -> Palette.to_bytes(palette)
+    end
+
     data =
       @signature <>
         Header.to_bytes(raw_data.header_chunk) <>
+        palette_data <>
         image_data <>
         End.to_bytes(raw_data.end_chunk)
 

lib/ex_png/utilities.ex

@@ -0,0 +1,16 @@
+defmodule ExPng.Utilities do
+  @moduledoc """
+  Shared utility functions.
+  """
+
+  @doc """
+  Accepts a list of binaries and reduces them to a single
+  binary
+  """
+  @spec reduce_to_binary([binary]) :: binary
+  def reduce_to_binary(list) do
+    list
+    |> Enum.reverse()
+    |> Enum.reduce(&Kernel.<>/2)
+  end
+end