{-# LANGUAGE ExistentialQuantification #-}
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE CPP #-}
-- | This module provide some color quantisation algorithm

-- in order to help in the creation of paletted images.

-- The most important function is `palettize` which will

-- make everything to create a nice color indexed image

-- with its palette.

module Codec.Picture.ColorQuant
    ( palettize
    , palettizeWithAlpha
    , defaultPaletteOptions
    , PaletteCreationMethod(..)
    , PaletteOptions( .. )
    ) where

#if !MIN_VERSION_base(4,8,0)
import           Control.Applicative (Applicative (..), (<$>))
#endif

import           Data.Bits           (unsafeShiftL, unsafeShiftR, (.&.), (.|.))
import           Data.List           (elemIndex)
import           Data.Maybe          (fromMaybe)
import           Data.Set            (Set)
import qualified Data.Set            as Set
import           Data.Word           (Word32)

import           Data.Vector         (Vector, (!))
import qualified Data.Vector         as V
import qualified Data.Vector.Unboxed as VU
import qualified Data.Vector.Storable as VS

import           Codec.Picture.Types
import           Codec.Picture.Gif (GifFrame(..), GifDisposalMethod, GifDelay)

-------------------------------------------------------------------------------

----            Palette Creation and Dithering

-------------------------------------------------------------------------------


-- | Define which palette creation method is used.

data PaletteCreationMethod =
      -- | MedianMeanCut method, provide the best results (visualy)

      -- at the cost of increased calculations.

      MedianMeanCut
      -- | Very fast algorithm (one pass), doesn't provide good

      -- looking results.

    | Uniform

-- | To specify how the palette will be created.

data PaletteOptions = PaletteOptions
    { -- | Algorithm used to find the palette

      PaletteOptions -> PaletteCreationMethod
paletteCreationMethod :: PaletteCreationMethod

      -- | Do we want to apply the dithering to the

      -- image. Enabling it often reduce compression

      -- ratio but enhance the perceived quality

      -- of the final image.

    , PaletteOptions -> Bool
enableImageDithering  :: Bool

      -- | Maximum number of color we want in the

      -- palette

    , PaletteOptions -> Int
paletteColorCount     :: Int
    }

-- | Default palette option, which aim at the best quality

-- and maximum possible colors (256)

defaultPaletteOptions :: PaletteOptions
defaultPaletteOptions :: PaletteOptions
defaultPaletteOptions = PaletteOptions :: PaletteCreationMethod -> Bool -> Int -> PaletteOptions
PaletteOptions
    { paletteCreationMethod :: PaletteCreationMethod
paletteCreationMethod = PaletteCreationMethod
MedianMeanCut
    , enableImageDithering :: Bool
enableImageDithering  = Bool
True
    , paletteColorCount :: Int
paletteColorCount     = Int
256
    }

-- | Changes all pixels with alpha = 0 to black

-- converting image to RGB (from RGBA) in meantime

alphaToBlack :: Image PixelRGBA8 -> Image PixelRGB8
alphaToBlack :: Image PixelRGBA8 -> Image PixelRGB8
alphaToBlack = (PixelRGBA8 -> PixelRGB8) -> Image PixelRGBA8 -> Image PixelRGB8
forall a b. (Pixel a, Pixel b) => (a -> b) -> Image a -> Image b
pixelMap PixelRGBA8 -> PixelRGB8
f
  where f :: PixelRGBA8 -> PixelRGB8
f (PixelRGBA8 Pixel8
r Pixel8
g Pixel8
b Pixel8
a) =
          if Pixel8
a Pixel8 -> Pixel8 -> Bool
forall a. Eq a => a -> a -> Bool
== Pixel8
0 then Pixel8 -> Pixel8 -> Pixel8 -> PixelRGB8
PixelRGB8 Pixel8
0 Pixel8
0 Pixel8
0
          else Pixel8 -> Pixel8 -> Pixel8 -> PixelRGB8
PixelRGB8 Pixel8
r Pixel8
g Pixel8
b

-- | Using second image as a stencil, changes palette index to the transparent

alphaTo255 :: Image Pixel8 -> Image PixelRGBA8 -> Pixel8 -> Image Pixel8
alphaTo255 :: Image Pixel8 -> Image PixelRGBA8 -> Pixel8 -> Image Pixel8
alphaTo255 Image Pixel8
img1 Image PixelRGBA8
img2 Pixel8
transparentIndex = (Int -> Int -> Pixel8) -> Int -> Int -> Image Pixel8
forall px. Pixel px => (Int -> Int -> px) -> Int -> Int -> Image px
generateImage Int -> Int -> Pixel8
f (Image Pixel8 -> Int
forall a. Image a -> Int
imageWidth Image Pixel8
img1) (Image PixelRGBA8 -> Int
forall a. Image a -> Int
imageHeight Image PixelRGBA8
img2)
  where f :: Int -> Int -> Pixel8
f Int
x Int
y =
          if Pixel8
a Pixel8 -> Pixel8 -> Bool
forall a. Eq a => a -> a -> Bool
== Pixel8
0 then Pixel8
transparentIndex
          else Pixel8
v
          where v :: Pixel8
v = Image Pixel8 -> Int -> Int -> Pixel8
forall a. Pixel a => Image a -> Int -> Int -> a
pixelAt Image Pixel8
img1 Int
x Int
y
                PixelRGBA8 Pixel8
_ Pixel8
_ Pixel8
_ Pixel8
a = Image PixelRGBA8 -> Int -> Int -> PixelRGBA8
forall a. Pixel a => Image a -> Int -> Int -> a
pixelAt Image PixelRGBA8
img2 Int
x Int
y

-- | Converts RGBA image to the array of GifFame's to use in encodeComplexGifImage

palettizeWithAlpha :: [(GifDelay, Image PixelRGBA8)] -> GifDisposalMethod -> [GifFrame]
palettizeWithAlpha :: [(Int, Image PixelRGBA8)] -> GifDisposalMethod -> [GifFrame]
palettizeWithAlpha [] GifDisposalMethod
_ = []
palettizeWithAlpha ((Int, Image PixelRGBA8)
x:[(Int, Image PixelRGBA8)]
xs) GifDisposalMethod
dispose =
  Int
-> Int
-> Maybe (Image PixelRGB8)
-> Maybe Int
-> Int
-> GifDisposalMethod
-> Image Pixel8
-> GifFrame
GifFrame
    Int
0 -- Offset X

    Int
0 -- Offset Y 

    (Image PixelRGB8 -> Maybe (Image PixelRGB8)
forall a. a -> Maybe a
Just (Image PixelRGB8 -> Maybe (Image PixelRGB8))
-> Image PixelRGB8 -> Maybe (Image PixelRGB8)
forall a b. (a -> b) -> a -> b
$ Image PixelRGB8
palet)
    (Int -> Maybe Int
forall a. a -> Maybe a
Just (Int -> Maybe Int) -> Int -> Maybe Int
forall a b. (a -> b) -> a -> b
$ Int
transparentIndex)
    Int
delay
    GifDisposalMethod
dispose
    (Image Pixel8 -> Image PixelRGBA8 -> Pixel8 -> Image Pixel8
alphaTo255 Image Pixel8
pixels Image PixelRGBA8
i (Int -> Pixel8
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
transparentIndex))
  GifFrame -> [GifFrame] -> [GifFrame]
forall a. a -> [a] -> [a]
: [(Int, Image PixelRGBA8)] -> GifDisposalMethod -> [GifFrame]
palettizeWithAlpha [(Int, Image PixelRGBA8)]
xs GifDisposalMethod
dispose
  where (Int
delay, Image PixelRGBA8
i) = (Int, Image PixelRGBA8)
x
        img :: Image PixelRGB8
img = Image PixelRGBA8 -> Image PixelRGB8
alphaToBlack Image PixelRGBA8
i
        (Image PixelRGB8
palet, Image Pixel8
pixels) =
          if Bool
isBelow
            then (Vector PixelRGB8 -> Image PixelRGB8
vecToPalette (Vector PixelRGB8
belowPaletteVec Vector PixelRGB8 -> PixelRGB8 -> Vector PixelRGB8
forall a. Vector a -> a -> Vector a
`V.snoc` Pixel8 -> Pixel8 -> Pixel8 -> PixelRGB8
PixelRGB8 Pixel8
0 Pixel8
0 Pixel8
0), (PixelRGB8 -> Pixel8) -> Image PixelRGB8 -> Image Pixel8
forall a b. (Pixel a, Pixel b) => (a -> b) -> Image a -> Image b
pixelMap PixelRGB8 -> Pixel8
belowPaletteIndex Image PixelRGB8
img)
            else (Vector PixelRGB8 -> Image PixelRGB8
vecToPalette (Vector PixelRGB8
genPaletteVec   Vector PixelRGB8 -> PixelRGB8 -> Vector PixelRGB8
forall a. Vector a -> a -> Vector a
`V.snoc` Pixel8 -> Pixel8 -> Pixel8 -> PixelRGB8
PixelRGB8 Pixel8
0 Pixel8
0 Pixel8
0), (PixelRGB8 -> Pixel8) -> Image PixelRGB8 -> Image Pixel8
forall a b. (Pixel a, Pixel b) => (a -> b) -> Image a -> Image b
pixelMap PixelRGB8 -> Pixel8
genPaletteIndex Image PixelRGB8
img)

        (Set PixelRGB8
belowPalette, Bool
isBelow) = Int -> Image PixelRGB8 -> (Set PixelRGB8, Bool)
isColorCountBelow Int
255 Image PixelRGB8
img
        belowPaletteVec :: Vector PixelRGB8
belowPaletteVec = [PixelRGB8] -> Vector PixelRGB8
forall a. [a] -> Vector a
V.fromList ([PixelRGB8] -> Vector PixelRGB8)
-> [PixelRGB8] -> Vector PixelRGB8
forall a b. (a -> b) -> a -> b
$ Set PixelRGB8 -> [PixelRGB8]
forall a. Set a -> [a]
Set.toList Set PixelRGB8
belowPalette
        belowPaletteIndex :: PixelRGB8 -> Pixel8
belowPaletteIndex PixelRGB8
p = PixelRGB8 -> Vector PixelRGB8 -> Pixel8
nearestColorIdx PixelRGB8
p Vector PixelRGB8
belowPaletteVec

        cs :: [Cluster]
cs = Set Cluster -> [Cluster]
forall a. Set a -> [a]
Set.toList (Set Cluster -> [Cluster])
-> (Image PixelRGB8 -> Set Cluster) -> Image PixelRGB8 -> [Cluster]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Image PixelRGB8 -> Set Cluster
clusters Int
255 (Image PixelRGB8 -> [Cluster]) -> Image PixelRGB8 -> [Cluster]
forall a b. (a -> b) -> a -> b
$ Image PixelRGB8
img
        genPaletteVec :: Vector PixelRGB8
genPaletteVec = [Cluster] -> Vector PixelRGB8
mkPaletteVec [Cluster]
cs
        genPaletteIndex :: PixelRGB8 -> Pixel8
genPaletteIndex PixelRGB8
p = PixelRGB8 -> Vector PixelRGB8 -> Pixel8
nearestColorIdx PixelRGB8
p Vector PixelRGB8
genPaletteVec

        transparentIndex :: Int
transparentIndex = Vector PixelRGB8 -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length (Vector PixelRGB8 -> Int) -> Vector PixelRGB8 -> Int
forall a b. (a -> b) -> a -> b
$ if Bool
isBelow then Vector PixelRGB8
belowPaletteVec else Vector PixelRGB8
genPaletteVec

-- | Reduces an image to a color palette according to `PaletteOptions` and

--   returns the /indices image/ along with its `Palette`.

palettize :: PaletteOptions -> Image PixelRGB8 -> (Image Pixel8, Palette)
palettize :: PaletteOptions
-> Image PixelRGB8 -> (Image Pixel8, Image PixelRGB8)
palettize opts :: PaletteOptions
opts@PaletteOptions { paletteCreationMethod :: PaletteOptions -> PaletteCreationMethod
paletteCreationMethod = PaletteCreationMethod
method } =
  case PaletteCreationMethod
method of
    PaletteCreationMethod
MedianMeanCut -> PaletteOptions
-> Image PixelRGB8 -> (Image Pixel8, Image PixelRGB8)
medianMeanCutQuantization PaletteOptions
opts
    PaletteCreationMethod
Uniform       -> PaletteOptions
-> Image PixelRGB8 -> (Image Pixel8, Image PixelRGB8)
uniformQuantization PaletteOptions
opts

-- | Modified median cut algorithm with optional ordered dithering. Returns an

-- image of `Pixel8` that acts as a matrix of indices into the `Palette`.

medianMeanCutQuantization :: PaletteOptions -> Image PixelRGB8
                          -> (Image Pixel8, Palette)
medianMeanCutQuantization :: PaletteOptions
-> Image PixelRGB8 -> (Image Pixel8, Image PixelRGB8)
medianMeanCutQuantization PaletteOptions
opts Image PixelRGB8
img
  | Bool
isBelow =
      ((PixelRGB8 -> Pixel8) -> Image PixelRGB8 -> Image Pixel8
forall a b. (Pixel a, Pixel b) => (a -> b) -> Image a -> Image b
pixelMap PixelRGB8 -> Pixel8
okPaletteIndex Image PixelRGB8
img, Vector PixelRGB8 -> Image PixelRGB8
vecToPalette Vector PixelRGB8
okPaletteVec)
  | PaletteOptions -> Bool
enableImageDithering PaletteOptions
opts = ((PixelRGB8 -> Pixel8) -> Image PixelRGB8 -> Image Pixel8
forall a b. (Pixel a, Pixel b) => (a -> b) -> Image a -> Image b
pixelMap PixelRGB8 -> Pixel8
paletteIndex Image PixelRGB8
dImg, Image PixelRGB8
palette)
  | Bool
otherwise = ((PixelRGB8 -> Pixel8) -> Image PixelRGB8 -> Image Pixel8
forall a b. (Pixel a, Pixel b) => (a -> b) -> Image a -> Image b
pixelMap PixelRGB8 -> Pixel8
paletteIndex Image PixelRGB8
img, Image PixelRGB8
palette)
  where
    maxColorCount :: Int
maxColorCount = PaletteOptions -> Int
paletteColorCount PaletteOptions
opts
    (Set PixelRGB8
okPalette, Bool
isBelow) = Int -> Image PixelRGB8 -> (Set PixelRGB8, Bool)
isColorCountBelow Int
maxColorCount Image PixelRGB8
img
    okPaletteVec :: Vector PixelRGB8
okPaletteVec = [PixelRGB8] -> Vector PixelRGB8
forall a. [a] -> Vector a
V.fromList ([PixelRGB8] -> Vector PixelRGB8)
-> [PixelRGB8] -> Vector PixelRGB8
forall a b. (a -> b) -> a -> b
$ Set PixelRGB8 -> [PixelRGB8]
forall a. Set a -> [a]
Set.toList Set PixelRGB8
okPalette
    okPaletteIndex :: PixelRGB8 -> Pixel8
okPaletteIndex PixelRGB8
p = PixelRGB8 -> Vector PixelRGB8 -> Pixel8
nearestColorIdx PixelRGB8
p Vector PixelRGB8
okPaletteVec

    palette :: Image PixelRGB8
palette = Vector PixelRGB8 -> Image PixelRGB8
vecToPalette Vector PixelRGB8
paletteVec
    paletteIndex :: PixelRGB8 -> Pixel8
paletteIndex PixelRGB8
p = PixelRGB8 -> Vector PixelRGB8 -> Pixel8
nearestColorIdx PixelRGB8
p Vector PixelRGB8
paletteVec
    paletteVec :: Vector PixelRGB8
paletteVec = [Cluster] -> Vector PixelRGB8
mkPaletteVec [Cluster]
cs
    cs :: [Cluster]
cs =  Set Cluster -> [Cluster]
forall a. Set a -> [a]
Set.toList (Set Cluster -> [Cluster])
-> (Image PixelRGB8 -> Set Cluster) -> Image PixelRGB8 -> [Cluster]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Image PixelRGB8 -> Set Cluster
clusters Int
maxColorCount (Image PixelRGB8 -> [Cluster]) -> Image PixelRGB8 -> [Cluster]
forall a b. (a -> b) -> a -> b
$ Image PixelRGB8
img
    dImg :: Image PixelRGB8
dImg = (Int -> Int -> PixelRGB8 -> PixelRGB8)
-> Image PixelRGB8 -> Image PixelRGB8
forall a b.
(Pixel a, Pixel b) =>
(Int -> Int -> a -> b) -> Image a -> Image b
pixelMapXY Int -> Int -> PixelRGB8 -> PixelRGB8
dither Image PixelRGB8
img

-- | A naive one pass Color Quantization algorithm - Uniform Quantization.

-- Simply take the most significant bits. The maxCols parameter is rounded

-- down to the nearest power of 2, and the bits are divided among the three

-- color channels with priority order green, red, blue. Returns an

-- image of `Pixel8` that acts as a matrix of indices into the `Palette`.

uniformQuantization :: PaletteOptions -> Image PixelRGB8 -> (Image Pixel8, Palette)
uniformQuantization :: PaletteOptions
-> Image PixelRGB8 -> (Image Pixel8, Image PixelRGB8)
uniformQuantization PaletteOptions
opts Image PixelRGB8
img
  -- -| colorCount img <= maxCols = colorQuantExact img

  | PaletteOptions -> Bool
enableImageDithering PaletteOptions
opts =
        ((PixelRGB8 -> Pixel8) -> Image PixelRGB8 -> Image Pixel8
forall a b. (Pixel a, Pixel b) => (a -> b) -> Image a -> Image b
pixelMap PixelRGB8 -> Pixel8
paletteIndex ((Int -> Int -> PixelRGB8 -> PixelRGB8)
-> Image PixelRGB8 -> Image PixelRGB8
forall a b.
(Pixel a, Pixel b) =>
(Int -> Int -> a -> b) -> Image a -> Image b
pixelMapXY Int -> Int -> PixelRGB8 -> PixelRGB8
dither Image PixelRGB8
img), Image PixelRGB8
palette)
  | Bool
otherwise = ((PixelRGB8 -> Pixel8) -> Image PixelRGB8 -> Image Pixel8
forall a b. (Pixel a, Pixel b) => (a -> b) -> Image a -> Image b
pixelMap PixelRGB8 -> Pixel8
paletteIndex Image PixelRGB8
img, Image PixelRGB8
palette)
  where
    maxCols :: Int
maxCols = PaletteOptions -> Int
paletteColorCount PaletteOptions
opts
    palette :: Image PixelRGB8
palette = [PixelRGB8] -> Image PixelRGB8
listToPalette [PixelRGB8]
paletteList
    paletteList :: [PixelRGB8]
paletteList = [Pixel8 -> Pixel8 -> Pixel8 -> PixelRGB8
PixelRGB8 Pixel8
r Pixel8
g Pixel8
b | Pixel8
r <- [Pixel8
0,Pixel8
dr..Pixel8
255]
                                   , Pixel8
g <- [Pixel8
0,Pixel8
dg..Pixel8
255]
                                   , Pixel8
b <- [Pixel8
0,Pixel8
db..Pixel8
255]]
    (Int
bg, Int
br, Int
bb) = Int -> (Int, Int, Int)
bitDiv3 Int
maxCols
    (Pixel8
dr, Pixel8
dg, Pixel8
db) = (Pixel8
2Pixel8 -> Int -> Pixel8
forall a b. (Num a, Integral b) => a -> b -> a
^(Int
8Int -> Int -> Int
forall a. Num a => a -> a -> a
-Int
br), Pixel8
2Pixel8 -> Int -> Pixel8
forall a b. (Num a, Integral b) => a -> b -> a
^(Int
8Int -> Int -> Int
forall a. Num a => a -> a -> a
-Int
bg), Pixel8
2Pixel8 -> Int -> Pixel8
forall a b. (Num a, Integral b) => a -> b -> a
^(Int
8Int -> Int -> Int
forall a. Num a => a -> a -> a
-Int
bb))
    paletteIndex :: PixelRGB8 -> Pixel8
paletteIndex (PixelRGB8 Pixel8
r Pixel8
g Pixel8
b) = Int -> Pixel8
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> Pixel8) -> Int -> Pixel8
forall a b. (a -> b) -> a -> b
$ Int -> Maybe Int -> Int
forall a. a -> Maybe a -> a
fromMaybe Int
0 (PixelRGB8 -> [PixelRGB8] -> Maybe Int
forall a. Eq a => a -> [a] -> Maybe Int
elemIndex
      (Pixel8 -> Pixel8 -> Pixel8 -> PixelRGB8
PixelRGB8 (Pixel8
r Pixel8 -> Pixel8 -> Pixel8
forall a. Bits a => a -> a -> a
.&. Pixel8 -> Pixel8
forall a. Num a => a -> a
negate Pixel8
dr) (Pixel8
g Pixel8 -> Pixel8 -> Pixel8
forall a. Bits a => a -> a -> a
.&. Pixel8 -> Pixel8
forall a. Num a => a -> a
negate Pixel8
dg) (Pixel8
b Pixel8 -> Pixel8 -> Pixel8
forall a. Bits a => a -> a -> a
.&. Pixel8 -> Pixel8
forall a. Num a => a -> a
negate Pixel8
db))
      [PixelRGB8]
paletteList)

isColorCountBelow :: Int -> Image PixelRGB8 -> (Set.Set PixelRGB8, Bool)
isColorCountBelow :: Int -> Image PixelRGB8 -> (Set PixelRGB8, Bool)
isColorCountBelow Int
maxColorCount Image PixelRGB8
img = Int -> Set PixelRGB8 -> (Set PixelRGB8, Bool)
go Int
0 Set PixelRGB8
forall a. Set a
Set.empty
  where rawData :: Vector (PixelBaseComponent PixelRGB8)
rawData = Image PixelRGB8 -> Vector (PixelBaseComponent PixelRGB8)
forall a. Image a -> Vector (PixelBaseComponent a)
imageData Image PixelRGB8
img
        maxIndex :: Int
maxIndex = Vector Pixel8 -> Int
forall a. Storable a => Vector a -> Int
VS.length Vector Pixel8
Vector (PixelBaseComponent PixelRGB8)
rawData
        
        go :: Int -> Set PixelRGB8 -> (Set PixelRGB8, Bool)
go !Int
idx !Set PixelRGB8
allColors
            | Set PixelRGB8 -> Int
forall a. Set a -> Int
Set.size Set PixelRGB8
allColors Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
maxColorCount = (Set PixelRGB8
forall a. Set a
Set.empty, Bool
False)
            | Int
idx Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
maxIndex Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
2 = (Set PixelRGB8
allColors, Bool
True)
            | Bool
otherwise = Int -> Set PixelRGB8 -> (Set PixelRGB8, Bool)
go (Int
idx Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
3) (Set PixelRGB8 -> (Set PixelRGB8, Bool))
-> Set PixelRGB8 -> (Set PixelRGB8, Bool)
forall a b. (a -> b) -> a -> b
$ PixelRGB8 -> Set PixelRGB8 -> Set PixelRGB8
forall a. Ord a => a -> Set a -> Set a
Set.insert PixelRGB8
px Set PixelRGB8
allColors
                where px :: PixelRGB8
px = Vector (PixelBaseComponent PixelRGB8) -> Int -> PixelRGB8
forall a. Pixel a => Vector (PixelBaseComponent a) -> Int -> a
unsafePixelAt Vector (PixelBaseComponent PixelRGB8)
rawData Int
idx 

vecToPalette :: Vector PixelRGB8 -> Palette
vecToPalette :: Vector PixelRGB8 -> Image PixelRGB8
vecToPalette Vector PixelRGB8
ps = (Int -> Int -> PixelRGB8) -> Int -> Int -> Image PixelRGB8
forall px. Pixel px => (Int -> Int -> px) -> Int -> Int -> Image px
generateImage (\Int
x Int
_ -> Vector PixelRGB8
ps Vector PixelRGB8 -> Int -> PixelRGB8
forall a. Vector a -> Int -> a
! Int
x) (Vector PixelRGB8 -> Int
forall a. Vector a -> Int
V.length Vector PixelRGB8
ps) Int
1

listToPalette :: [PixelRGB8] -> Palette
listToPalette :: [PixelRGB8] -> Image PixelRGB8
listToPalette [PixelRGB8]
ps = (Int -> Int -> PixelRGB8) -> Int -> Int -> Image PixelRGB8
forall px. Pixel px => (Int -> Int -> px) -> Int -> Int -> Image px
generateImage (\Int
x Int
_ -> [PixelRGB8]
ps [PixelRGB8] -> Int -> PixelRGB8
forall a. [a] -> Int -> a
!! Int
x) ([PixelRGB8] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [PixelRGB8]
ps) Int
1

bitDiv3 :: Int -> (Int, Int, Int)
bitDiv3 :: Int -> (Int, Int, Int)
bitDiv3 Int
n = case Int
r of
            Int
0 -> (Int
q, Int
q, Int
q)
            Int
1 -> (Int
qInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1, Int
q, Int
q)
            Int
_ -> (Int
qInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1, Int
qInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1, Int
q)
  where
    r :: Int
r = Int
m Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` Int
3
    q :: Int
q = Int
m Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
3
    m :: Int
m = Double -> Int
forall a b. (RealFrac a, Integral b) => a -> b
floor (Double -> Int) -> (Double -> Double) -> Double -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Double -> Double -> Double
forall a. Floating a => a -> a -> a
logBase (Double
2 :: Double) (Double -> Int) -> Double -> Int
forall a b. (a -> b) -> a -> b
$ Int -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
n

-------------------------------------------------------------------------------

----            Dithering

-------------------------------------------------------------------------------


-- Add a dither mask to an image for ordered dithering.

-- Uses a small, spatially stable dithering algorithm based on magic numbers

-- and arithmetic inspired by the /a dither/ algorithm of Øyvind Kolås,

-- pippin@gimp.org, 2013. See, http://pippin.gimp.org/a_dither/.

dither :: Int -> Int -> PixelRGB8 -> PixelRGB8
dither :: Int -> Int -> PixelRGB8 -> PixelRGB8
dither Int
x Int
y (PixelRGB8 Pixel8
r Pixel8
g Pixel8
b) = Pixel8 -> Pixel8 -> Pixel8 -> PixelRGB8
PixelRGB8 (Int -> Pixel8
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
r')
                                         (Int -> Pixel8
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
g')
                                         (Int -> Pixel8
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
b')
  where
    -- Should view 16 as a parameter that can be optimized for best looking

    -- results

    r' :: Int
r' = Int -> Int -> Int
forall a. Ord a => a -> a -> a
min Int
255 (Pixel8 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Pixel8
r Int -> Int -> Int
forall a. Num a => a -> a -> a
+ (Int
x' Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
y') Int -> Int -> Int
forall a. Bits a => a -> a -> a
.&. Int
16)
    g' :: Int
g' = Int -> Int -> Int
forall a. Ord a => a -> a -> a
min Int
255 (Pixel8 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Pixel8
g Int -> Int -> Int
forall a. Num a => a -> a -> a
+ (Int
x' Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
y' Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
7973) Int -> Int -> Int
forall a. Bits a => a -> a -> a
.&. Int
16)
    b' :: Int
b' = Int -> Int -> Int
forall a. Ord a => a -> a -> a
min Int
255 (Pixel8 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Pixel8
b Int -> Int -> Int
forall a. Num a => a -> a -> a
+ (Int
x' Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
y' Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
15946) Int -> Int -> Int
forall a. Bits a => a -> a -> a
.&. Int
16)
    x' :: Int
x' = Int
119 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
x
    y' :: Int
y' = Int
28084 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
y

-------------------------------------------------------------------------------

----            Small modification of foldl package by Gabriel Gonzalez

-------------------------------------------------------------------------------


-- Modification to Control.foldl by Gabriel Gonzalez copyright 2013, BSD3.

-- http://hackage.haskell.org/package/foldl-1.0.1/docs/Control-Foldl.html


{-| Efficient representation of a left fold that preserves the fold's step
    function, initial accumulator, and extraction function

    This allows the 'Applicative' instance to assemble derived folds that
    traverse the container only once
-}
data Fold a b = forall x . Fold (x -> a -> x) x (x -> b)

{-| Apply a strict left 'Fold' to a 'Foldable' container

    Much slower than 'fold' on lists because 'Foldable' operations currently do
    not trigger @build/foldr@ fusion
-}
fold :: Fold PackedRGB b -> VU.Vector PackedRGB -> b
fold :: Fold PackedRGB b -> Vector PackedRGB -> b
fold (Fold x -> PackedRGB -> x
step x
begin x -> b
done) = x -> b
done (x -> b) -> (Vector PackedRGB -> x) -> Vector PackedRGB -> b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (x -> PackedRGB -> x) -> x -> Vector PackedRGB -> x
forall b a. Unbox b => (a -> b -> a) -> a -> Vector b -> a
VU.foldl' x -> PackedRGB -> x
step x
begin
{-# INLINE fold #-}

{-
F.foldr :: (a -> b -> b) -> b -> t a -> b

fold :: (Foldable f) => Fold a b -> f a -> b
fold (Fold step begin done) as = F.foldr step' done as begin
  where step' x k z = k $! step z x
-}

data Pair a b = Pair !a !b

instance Functor (Fold a) where
    fmap :: (a -> b) -> Fold a a -> Fold a b
fmap a -> b
f (Fold x -> a -> x
step x
begin x -> a
done) = (x -> a -> x) -> x -> (x -> b) -> Fold a b
forall a b x. (x -> a -> x) -> x -> (x -> b) -> Fold a b
Fold x -> a -> x
step x
begin (a -> b
f (a -> b) -> (x -> a) -> x -> b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. x -> a
done)
    {-# INLINABLE fmap #-}

instance Applicative (Fold a) where
    pure :: a -> Fold a a
pure a
b    = (() -> a -> ()) -> () -> (() -> a) -> Fold a a
forall a b x. (x -> a -> x) -> x -> (x -> b) -> Fold a b
Fold (\() a
_ -> ()) () (\() -> a
b)
    {-# INLINABLE pure #-}
    (Fold x -> a -> x
stepL x
beginL x -> a -> b
doneL) <*> :: Fold a (a -> b) -> Fold a a -> Fold a b
<*> (Fold x -> a -> x
stepR x
beginR x -> a
doneR) =
        let step :: Pair x x -> a -> Pair x x
step (Pair x
xL x
xR) a
a = x -> x -> Pair x x
forall a b. a -> b -> Pair a b
Pair (x -> a -> x
stepL x
xL a
a) (x -> a -> x
stepR x
xR a
a)
            begin :: Pair x x
begin = x -> x -> Pair x x
forall a b. a -> b -> Pair a b
Pair x
beginL x
beginR
            done :: Pair x x -> b
done (Pair x
xL x
xR) = x -> a -> b
doneL x
xL (a -> b) -> a -> b
forall a b. (a -> b) -> a -> b
$ x -> a
doneR x
xR
        in  (Pair x x -> a -> Pair x x)
-> Pair x x -> (Pair x x -> b) -> Fold a b
forall a b x. (x -> a -> x) -> x -> (x -> b) -> Fold a b
Fold Pair x x -> a -> Pair x x
step Pair x x
begin Pair x x -> b
done
    {-# INLINABLE (<*>) #-}

{- | Like 'length', except with a more general 'Num' return value -}
intLength :: Fold a Int
intLength :: Fold a Int
intLength = (Int -> a -> Int) -> Int -> (Int -> Int) -> Fold a Int
forall a b x. (x -> a -> x) -> x -> (x -> b) -> Fold a b
Fold (\Int
n a
_ -> Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Int
0 Int -> Int
forall a. a -> a
id

-------------------------------------------------------------------------------

----            Modified Median Cut Algorithm

-------------------------------------------------------------------------------


-- Based on the OCaml implementation:

-- http://rosettacode.org/wiki/Color_quantization

-- which is in turn based on: www.leptonica.com/papers/mediancut.pdf.

-- We use the product of volume and population to determine the next cluster

-- to split and determine the placement of each color by compating it to the

-- mean of the parent cluster. So median cut is a bit of a misnomer, since one

-- of the modifiations is to use the mean.


mkPaletteVec :: [Cluster] -> Vector PixelRGB8
mkPaletteVec :: [Cluster] -> Vector PixelRGB8
mkPaletteVec  = [PixelRGB8] -> Vector PixelRGB8
forall a. [a] -> Vector a
V.fromList ([PixelRGB8] -> Vector PixelRGB8)
-> ([Cluster] -> [PixelRGB8]) -> [Cluster] -> Vector PixelRGB8
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Cluster -> PixelRGB8) -> [Cluster] -> [PixelRGB8]
forall a b. (a -> b) -> [a] -> [b]
map (PixelRGBF -> PixelRGB8
toRGB8 (PixelRGBF -> PixelRGB8)
-> (Cluster -> PixelRGBF) -> Cluster -> PixelRGB8
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Cluster -> PixelRGBF
meanColor)

type PackedRGB = Word32

data Cluster = Cluster
    { Cluster -> Float
value       :: {-# UNPACK #-} !Float
    , Cluster -> PixelRGBF
meanColor   :: !PixelRGBF
    , Cluster -> PixelRGBF
dims        :: !PixelRGBF
    , Cluster -> Vector PackedRGB
colors      :: VU.Vector PackedRGB
    }

instance Eq Cluster where
    Cluster
a == :: Cluster -> Cluster -> Bool
== Cluster
b =
        (Cluster -> Float
value Cluster
a, Cluster -> PixelRGBF
meanColor Cluster
a, Cluster -> PixelRGBF
dims Cluster
a) (Float, PixelRGBF, PixelRGBF)
-> (Float, PixelRGBF, PixelRGBF) -> Bool
forall a. Eq a => a -> a -> Bool
== (Cluster -> Float
value Cluster
b, Cluster -> PixelRGBF
meanColor Cluster
b, Cluster -> PixelRGBF
dims Cluster
b)

instance Ord Cluster where
    compare :: Cluster -> Cluster -> Ordering
compare Cluster
a Cluster
b =
        (Float, PixelRGBF, PixelRGBF)
-> (Float, PixelRGBF, PixelRGBF) -> Ordering
forall a. Ord a => a -> a -> Ordering
compare (Cluster -> Float
value Cluster
a, Cluster -> PixelRGBF
meanColor Cluster
a, Cluster -> PixelRGBF
dims Cluster
a) (Cluster -> Float
value Cluster
b, Cluster -> PixelRGBF
meanColor Cluster
b, Cluster -> PixelRGBF
dims Cluster
b)

data Axis = RAxis | GAxis | BAxis

inf :: Float
inf :: Float
inf = String -> Float
forall a. Read a => String -> a
read String
"Infinity"

fromRGB8 :: PixelRGB8 -> PixelRGBF
fromRGB8 :: PixelRGB8 -> PixelRGBF
fromRGB8 (PixelRGB8 Pixel8
r Pixel8
g Pixel8
b) =
  Float -> Float -> Float -> PixelRGBF
PixelRGBF (Pixel8 -> Float
forall a b. (Integral a, Num b) => a -> b
fromIntegral Pixel8
r) (Pixel8 -> Float
forall a b. (Integral a, Num b) => a -> b
fromIntegral Pixel8
g) (Pixel8 -> Float
forall a b. (Integral a, Num b) => a -> b
fromIntegral Pixel8
b)

toRGB8 :: PixelRGBF -> PixelRGB8
toRGB8 :: PixelRGBF -> PixelRGB8
toRGB8 (PixelRGBF Float
r Float
g Float
b) =
  Pixel8 -> Pixel8 -> Pixel8 -> PixelRGB8
PixelRGB8 (Float -> Pixel8
forall a b. (RealFrac a, Integral b) => a -> b
round Float
r) (Float -> Pixel8
forall a b. (RealFrac a, Integral b) => a -> b
round Float
g) (Float -> Pixel8
forall a b. (RealFrac a, Integral b) => a -> b
round Float
b)

meanRGB :: Fold PixelRGBF PixelRGBF
meanRGB :: Fold PixelRGBF PixelRGBF
meanRGB = Int -> PixelRGBF -> PixelRGBF
forall a a.
(Integral a, Pixel a, Fractional (PixelBaseComponent a)) =>
a -> a -> a
mean (Int -> PixelRGBF -> PixelRGBF)
-> Fold PixelRGBF Int -> Fold PixelRGBF (PixelRGBF -> PixelRGBF)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Fold PixelRGBF Int
forall a. Fold a Int
intLength Fold PixelRGBF (PixelRGBF -> PixelRGBF)
-> Fold PixelRGBF PixelRGBF -> Fold PixelRGBF PixelRGBF
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Fold PixelRGBF PixelRGBF
pixelSum
  where
    pixelSum :: Fold PixelRGBF PixelRGBF
pixelSum = (PixelRGBF -> PixelRGBF -> PixelRGBF)
-> PixelRGBF
-> (PixelRGBF -> PixelRGBF)
-> Fold PixelRGBF PixelRGBF
forall a b x. (x -> a -> x) -> x -> (x -> b) -> Fold a b
Fold ((Int
 -> PixelBaseComponent PixelRGBF
 -> PixelBaseComponent PixelRGBF
 -> PixelBaseComponent PixelRGBF)
-> PixelRGBF -> PixelRGBF -> PixelRGBF
forall a.
Pixel a =>
(Int
 -> PixelBaseComponent a
 -> PixelBaseComponent a
 -> PixelBaseComponent a)
-> a -> a -> a
mixWith ((Int
  -> PixelBaseComponent PixelRGBF
  -> PixelBaseComponent PixelRGBF
  -> PixelBaseComponent PixelRGBF)
 -> PixelRGBF -> PixelRGBF -> PixelRGBF)
-> (Int
    -> PixelBaseComponent PixelRGBF
    -> PixelBaseComponent PixelRGBF
    -> PixelBaseComponent PixelRGBF)
-> PixelRGBF
-> PixelRGBF
-> PixelRGBF
forall a b. (a -> b) -> a -> b
$ (Float -> Float -> Float) -> Int -> Float -> Float -> Float
forall a b. a -> b -> a
const Float -> Float -> Float
forall a. Num a => a -> a -> a
(+)) (Float -> Float -> Float -> PixelRGBF
PixelRGBF Float
0 Float
0 Float
0) PixelRGBF -> PixelRGBF
forall a. a -> a
id
    mean :: a -> a -> a
mean a
n = (PixelBaseComponent a -> PixelBaseComponent a) -> a -> a
forall a.
Pixel a =>
(PixelBaseComponent a -> PixelBaseComponent a) -> a -> a
colorMap (PixelBaseComponent a
-> PixelBaseComponent a -> PixelBaseComponent a
forall a. Fractional a => a -> a -> a
/ PixelBaseComponent a
nf)
      where nf :: PixelBaseComponent a
nf = a -> PixelBaseComponent a
forall a b. (Integral a, Num b) => a -> b
fromIntegral a
n

minimal :: Fold PixelRGBF PixelRGBF
minimal :: Fold PixelRGBF PixelRGBF
minimal = (PixelRGBF -> PixelRGBF -> PixelRGBF)
-> PixelRGBF
-> (PixelRGBF -> PixelRGBF)
-> Fold PixelRGBF PixelRGBF
forall a b x. (x -> a -> x) -> x -> (x -> b) -> Fold a b
Fold PixelRGBF -> PixelRGBF -> PixelRGBF
mini (Float -> Float -> Float -> PixelRGBF
PixelRGBF Float
inf Float
inf Float
inf) PixelRGBF -> PixelRGBF
forall a. a -> a
id
  where mini :: PixelRGBF -> PixelRGBF -> PixelRGBF
mini = (Int
 -> PixelBaseComponent PixelRGBF
 -> PixelBaseComponent PixelRGBF
 -> PixelBaseComponent PixelRGBF)
-> PixelRGBF -> PixelRGBF -> PixelRGBF
forall a.
Pixel a =>
(Int
 -> PixelBaseComponent a
 -> PixelBaseComponent a
 -> PixelBaseComponent a)
-> a -> a -> a
mixWith ((Int
  -> PixelBaseComponent PixelRGBF
  -> PixelBaseComponent PixelRGBF
  -> PixelBaseComponent PixelRGBF)
 -> PixelRGBF -> PixelRGBF -> PixelRGBF)
-> (Int
    -> PixelBaseComponent PixelRGBF
    -> PixelBaseComponent PixelRGBF
    -> PixelBaseComponent PixelRGBF)
-> PixelRGBF
-> PixelRGBF
-> PixelRGBF
forall a b. (a -> b) -> a -> b
$ (PixelBaseComponent PixelRGBF
 -> PixelBaseComponent PixelRGBF -> PixelBaseComponent PixelRGBF)
-> Int
-> PixelBaseComponent PixelRGBF
-> PixelBaseComponent PixelRGBF
-> PixelBaseComponent PixelRGBF
forall a b. a -> b -> a
const PixelBaseComponent PixelRGBF
-> PixelBaseComponent PixelRGBF -> PixelBaseComponent PixelRGBF
forall a. Ord a => a -> a -> a
min

maximal :: Fold PixelRGBF PixelRGBF
maximal :: Fold PixelRGBF PixelRGBF
maximal = (PixelRGBF -> PixelRGBF -> PixelRGBF)
-> PixelRGBF
-> (PixelRGBF -> PixelRGBF)
-> Fold PixelRGBF PixelRGBF
forall a b x. (x -> a -> x) -> x -> (x -> b) -> Fold a b
Fold PixelRGBF -> PixelRGBF -> PixelRGBF
maxi (Float -> Float -> Float -> PixelRGBF
PixelRGBF (-Float
inf) (-Float
inf) (-Float
inf)) PixelRGBF -> PixelRGBF
forall a. a -> a
id
  where maxi :: PixelRGBF -> PixelRGBF -> PixelRGBF
maxi = (Int
 -> PixelBaseComponent PixelRGBF
 -> PixelBaseComponent PixelRGBF
 -> PixelBaseComponent PixelRGBF)
-> PixelRGBF -> PixelRGBF -> PixelRGBF
forall a.
Pixel a =>
(Int
 -> PixelBaseComponent a
 -> PixelBaseComponent a
 -> PixelBaseComponent a)
-> a -> a -> a
mixWith ((Int
  -> PixelBaseComponent PixelRGBF
  -> PixelBaseComponent PixelRGBF
  -> PixelBaseComponent PixelRGBF)
 -> PixelRGBF -> PixelRGBF -> PixelRGBF)
-> (Int
    -> PixelBaseComponent PixelRGBF
    -> PixelBaseComponent PixelRGBF
    -> PixelBaseComponent PixelRGBF)
-> PixelRGBF
-> PixelRGBF
-> PixelRGBF
forall a b. (a -> b) -> a -> b
$ (PixelBaseComponent PixelRGBF
 -> PixelBaseComponent PixelRGBF -> PixelBaseComponent PixelRGBF)
-> Int
-> PixelBaseComponent PixelRGBF
-> PixelBaseComponent PixelRGBF
-> PixelBaseComponent PixelRGBF
forall a b. a -> b -> a
const PixelBaseComponent PixelRGBF
-> PixelBaseComponent PixelRGBF -> PixelBaseComponent PixelRGBF
forall a. Ord a => a -> a -> a
max

extrems :: Fold PixelRGBF (PixelRGBF, PixelRGBF)
extrems :: Fold PixelRGBF (PixelRGBF, PixelRGBF)
extrems = (,) (PixelRGBF -> PixelRGBF -> (PixelRGBF, PixelRGBF))
-> Fold PixelRGBF PixelRGBF
-> Fold PixelRGBF (PixelRGBF -> (PixelRGBF, PixelRGBF))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Fold PixelRGBF PixelRGBF
minimal Fold PixelRGBF (PixelRGBF -> (PixelRGBF, PixelRGBF))
-> Fold PixelRGBF PixelRGBF
-> Fold PixelRGBF (PixelRGBF, PixelRGBF)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Fold PixelRGBF PixelRGBF
maximal

volAndDims :: Fold PixelRGBF (Float, PixelRGBF)
volAndDims :: Fold PixelRGBF (Float, PixelRGBF)
volAndDims = (PixelRGBF, PixelRGBF) -> (Float, PixelRGBF)
deltify ((PixelRGBF, PixelRGBF) -> (Float, PixelRGBF))
-> Fold PixelRGBF (PixelRGBF, PixelRGBF)
-> Fold PixelRGBF (Float, PixelRGBF)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Fold PixelRGBF (PixelRGBF, PixelRGBF)
extrems
  where deltify :: (PixelRGBF, PixelRGBF) -> (Float, PixelRGBF)
deltify (PixelRGBF
mini, PixelRGBF
maxi) = (Float
dr Float -> Float -> Float
forall a. Num a => a -> a -> a
* Float
dg Float -> Float -> Float
forall a. Num a => a -> a -> a
* Float
db, PixelRGBF
delta)
          where delta :: PixelRGBF
delta@(PixelRGBF Float
dr Float
dg Float
db) =
                        (Int
 -> PixelBaseComponent PixelRGBF
 -> PixelBaseComponent PixelRGBF
 -> PixelBaseComponent PixelRGBF)
-> PixelRGBF -> PixelRGBF -> PixelRGBF
forall a.
Pixel a =>
(Int
 -> PixelBaseComponent a
 -> PixelBaseComponent a
 -> PixelBaseComponent a)
-> a -> a -> a
mixWith ((Float -> Float -> Float) -> Int -> Float -> Float -> Float
forall a b. a -> b -> a
const (-)) PixelRGBF
maxi PixelRGBF
mini

unpackFold :: Fold PixelRGBF a -> Fold PackedRGB a
unpackFold :: Fold PixelRGBF a -> Fold PackedRGB a
unpackFold (Fold x -> PixelRGBF -> x
step x
start x -> a
done) = (x -> PackedRGB -> x) -> x -> (x -> a) -> Fold PackedRGB a
forall a b x. (x -> a -> x) -> x -> (x -> b) -> Fold a b
Fold (\x
acc -> x -> PixelRGBF -> x
step x
acc (PixelRGBF -> x) -> (PackedRGB -> PixelRGBF) -> PackedRGB -> x
forall b c a. (b -> c) -> (a -> b) -> a -> c
. PackedRGB -> PixelRGBF
transform) x
start x -> a
done
  where transform :: PackedRGB -> PixelRGBF
transform = PixelRGB8 -> PixelRGBF
fromRGB8 (PixelRGB8 -> PixelRGBF)
-> (PackedRGB -> PixelRGB8) -> PackedRGB -> PixelRGBF
forall b c a. (b -> c) -> (a -> b) -> a -> c
. PackedRGB -> PixelRGB8
rgbIntUnpack

mkCluster :: VU.Vector PackedRGB -> Cluster
mkCluster :: Vector PackedRGB -> Cluster
mkCluster Vector PackedRGB
ps = Cluster :: Float -> PixelRGBF -> PixelRGBF -> Vector PackedRGB -> Cluster
Cluster
    { value :: Float
value = Float
v Float -> Float -> Float
forall a. Num a => a -> a -> a
* Int -> Float
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
l
    , meanColor :: PixelRGBF
meanColor = PixelRGBF
m
    , dims :: PixelRGBF
dims = PixelRGBF
ds
    , colors :: Vector PackedRGB
colors = Vector PackedRGB
ps
    }
  where
    worker :: Fold PixelRGBF ((Float, PixelRGBF), PixelRGBF, Int)
worker = (,,) ((Float, PixelRGBF)
 -> PixelRGBF -> Int -> ((Float, PixelRGBF), PixelRGBF, Int))
-> Fold PixelRGBF (Float, PixelRGBF)
-> Fold
     PixelRGBF
     (PixelRGBF -> Int -> ((Float, PixelRGBF), PixelRGBF, Int))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Fold PixelRGBF (Float, PixelRGBF)
volAndDims Fold
  PixelRGBF
  (PixelRGBF -> Int -> ((Float, PixelRGBF), PixelRGBF, Int))
-> Fold PixelRGBF PixelRGBF
-> Fold PixelRGBF (Int -> ((Float, PixelRGBF), PixelRGBF, Int))
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Fold PixelRGBF PixelRGBF
meanRGB Fold PixelRGBF (Int -> ((Float, PixelRGBF), PixelRGBF, Int))
-> Fold PixelRGBF Int
-> Fold PixelRGBF ((Float, PixelRGBF), PixelRGBF, Int)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Fold PixelRGBF Int
forall a. Fold a Int
intLength
    ((Float
v, PixelRGBF
ds), PixelRGBF
m, Int
l) = Fold PackedRGB ((Float, PixelRGBF), PixelRGBF, Int)
-> Vector PackedRGB -> ((Float, PixelRGBF), PixelRGBF, Int)
forall b. Fold PackedRGB b -> Vector PackedRGB -> b
fold (Fold PixelRGBF ((Float, PixelRGBF), PixelRGBF, Int)
-> Fold PackedRGB ((Float, PixelRGBF), PixelRGBF, Int)
forall a. Fold PixelRGBF a -> Fold PackedRGB a
unpackFold Fold PixelRGBF ((Float, PixelRGBF), PixelRGBF, Int)
worker) Vector PackedRGB
ps

maxAxis :: PixelRGBF -> Axis
maxAxis :: PixelRGBF -> Axis
maxAxis (PixelRGBF Float
r Float
g Float
b) =
  case (Float
r Float -> Float -> Ordering
forall a. Ord a => a -> a -> Ordering
`compare` Float
g, Float
r Float -> Float -> Ordering
forall a. Ord a => a -> a -> Ordering
`compare` Float
b, Float
g Float -> Float -> Ordering
forall a. Ord a => a -> a -> Ordering
`compare` Float
b) of
    (Ordering
GT, Ordering
GT, Ordering
_)  -> Axis
RAxis
    (Ordering
LT, Ordering
GT, Ordering
_)  -> Axis
GAxis
    (Ordering
GT, Ordering
LT, Ordering
_)  -> Axis
BAxis
    (Ordering
LT, Ordering
LT, Ordering
GT) -> Axis
GAxis
    (Ordering
EQ, Ordering
GT, Ordering
_)  -> Axis
RAxis
    (Ordering
_,  Ordering
_,  Ordering
_)  -> Axis
BAxis

-- Split a cluster about its largest axis using the mean to divide up the

-- pixels.

subdivide :: Cluster -> (Cluster, Cluster)
subdivide :: Cluster -> (Cluster, Cluster)
subdivide Cluster
cluster = (Vector PackedRGB -> Cluster
mkCluster Vector PackedRGB
px1, Vector PackedRGB -> Cluster
mkCluster Vector PackedRGB
px2)
  where
    (PixelRGBF Float
mr Float
mg Float
mb) = Cluster -> PixelRGBF
meanColor Cluster
cluster
    (Vector PackedRGB
px1, Vector PackedRGB
px2) = (PackedRGB -> Bool)
-> Vector PackedRGB -> (Vector PackedRGB, Vector PackedRGB)
forall a.
Unbox a =>
(a -> Bool) -> Vector a -> (Vector a, Vector a)
VU.partition (PixelRGB8 -> Bool
cond (PixelRGB8 -> Bool)
-> (PackedRGB -> PixelRGB8) -> PackedRGB -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. PackedRGB -> PixelRGB8
rgbIntUnpack) (Vector PackedRGB -> (Vector PackedRGB, Vector PackedRGB))
-> Vector PackedRGB -> (Vector PackedRGB, Vector PackedRGB)
forall a b. (a -> b) -> a -> b
$ Cluster -> Vector PackedRGB
colors Cluster
cluster
    cond :: PixelRGB8 -> Bool
cond = case PixelRGBF -> Axis
maxAxis (PixelRGBF -> Axis) -> PixelRGBF -> Axis
forall a b. (a -> b) -> a -> b
$ Cluster -> PixelRGBF
dims Cluster
cluster of
      Axis
RAxis -> \(PixelRGB8 Pixel8
r Pixel8
_ Pixel8
_) -> Pixel8 -> Float
forall a b. (Integral a, Num b) => a -> b
fromIntegral Pixel8
r Float -> Float -> Bool
forall a. Ord a => a -> a -> Bool
< Float
mr
      Axis
GAxis -> \(PixelRGB8 Pixel8
_ Pixel8
g Pixel8
_) -> Pixel8 -> Float
forall a b. (Integral a, Num b) => a -> b
fromIntegral Pixel8
g Float -> Float -> Bool
forall a. Ord a => a -> a -> Bool
< Float
mg
      Axis
BAxis -> \(PixelRGB8 Pixel8
_ Pixel8
_ Pixel8
b) -> Pixel8 -> Float
forall a b. (Integral a, Num b) => a -> b
fromIntegral Pixel8
b Float -> Float -> Bool
forall a. Ord a => a -> a -> Bool
< Float
mb

rgbIntPack :: PixelRGB8 -> PackedRGB
rgbIntPack :: PixelRGB8 -> PackedRGB
rgbIntPack (PixelRGB8 Pixel8
r Pixel8
g Pixel8
b) =
    PackedRGB
wr PackedRGB -> Int -> PackedRGB
forall a. Bits a => a -> Int -> a
`unsafeShiftL` (Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
8) PackedRGB -> PackedRGB -> PackedRGB
forall a. Bits a => a -> a -> a
.|. PackedRGB
wg PackedRGB -> Int -> PackedRGB
forall a. Bits a => a -> Int -> a
`unsafeShiftL` Int
8 PackedRGB -> PackedRGB -> PackedRGB
forall a. Bits a => a -> a -> a
.|. PackedRGB
wb
  where wr :: PackedRGB
wr = Pixel8 -> PackedRGB
forall a b. (Integral a, Num b) => a -> b
fromIntegral Pixel8
r
        wg :: PackedRGB
wg = Pixel8 -> PackedRGB
forall a b. (Integral a, Num b) => a -> b
fromIntegral Pixel8
g
        wb :: PackedRGB
wb = Pixel8 -> PackedRGB
forall a b. (Integral a, Num b) => a -> b
fromIntegral Pixel8
b

rgbIntUnpack :: PackedRGB -> PixelRGB8
rgbIntUnpack :: PackedRGB -> PixelRGB8
rgbIntUnpack PackedRGB
v = Pixel8 -> Pixel8 -> Pixel8 -> PixelRGB8
PixelRGB8 Pixel8
r Pixel8
g Pixel8
b
  where
    r :: Pixel8
r = PackedRGB -> Pixel8
forall a b. (Integral a, Num b) => a -> b
fromIntegral (PackedRGB -> Pixel8) -> PackedRGB -> Pixel8
forall a b. (a -> b) -> a -> b
$ PackedRGB
v PackedRGB -> Int -> PackedRGB
forall a. Bits a => a -> Int -> a
`unsafeShiftR` (Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
8)
    g :: Pixel8
g = PackedRGB -> Pixel8
forall a b. (Integral a, Num b) => a -> b
fromIntegral (PackedRGB -> Pixel8) -> PackedRGB -> Pixel8
forall a b. (a -> b) -> a -> b
$ PackedRGB
v PackedRGB -> Int -> PackedRGB
forall a. Bits a => a -> Int -> a
`unsafeShiftR` Int
8
    b :: Pixel8
b = PackedRGB -> Pixel8
forall a b. (Integral a, Num b) => a -> b
fromIntegral PackedRGB
v

initCluster :: Image PixelRGB8 -> Cluster
initCluster :: Image PixelRGB8 -> Cluster
initCluster Image PixelRGB8
img = Vector PackedRGB -> Cluster
mkCluster (Vector PackedRGB -> Cluster) -> Vector PackedRGB -> Cluster
forall a b. (a -> b) -> a -> b
$ Int -> (Int -> PackedRGB) -> Vector PackedRGB
forall a. Unbox a => Int -> (Int -> a) -> Vector a
VU.generate ((Int
w Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
h) Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
subSampling) Int -> PackedRGB
packer
  where samplingFactor :: Int
samplingFactor = Int
3
        subSampling :: Int
subSampling = Int
samplingFactor Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
samplingFactor
        compCount :: Int
compCount = PixelRGB8 -> Int
forall a. Pixel a => a -> Int
componentCount (PixelRGB8
forall a. HasCallStack => a
undefined :: PixelRGB8)
        w :: Int
w = Image PixelRGB8 -> Int
forall a. Image a -> Int
imageWidth Image PixelRGB8
img
        h :: Int
h = Image PixelRGB8 -> Int
forall a. Image a -> Int
imageHeight Image PixelRGB8
img
        rawData :: Vector (PixelBaseComponent PixelRGB8)
rawData = Image PixelRGB8 -> Vector (PixelBaseComponent PixelRGB8)
forall a. Image a -> Vector (PixelBaseComponent a)
imageData Image PixelRGB8
img
        packer :: Int -> PackedRGB
packer Int
ix =
            PixelRGB8 -> PackedRGB
rgbIntPack (PixelRGB8 -> PackedRGB) -> (Int -> PixelRGB8) -> Int -> PackedRGB
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Vector (PixelBaseComponent PixelRGB8) -> Int -> PixelRGB8
forall a. Pixel a => Vector (PixelBaseComponent a) -> Int -> a
unsafePixelAt Vector (PixelBaseComponent PixelRGB8)
rawData (Int -> PackedRGB) -> Int -> PackedRGB
forall a b. (a -> b) -> a -> b
$ Int
ix Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
subSampling Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
compCount

-- Take the cluster with the largest value = (volume * population) and remove it

-- from the priority queue. Then subdivide it about its largest axis and put the

-- two new clusters on the queue.

split :: Set Cluster -> Set Cluster
split :: Set Cluster -> Set Cluster
split Set Cluster
cs = Cluster -> Set Cluster -> Set Cluster
forall a. Ord a => a -> Set a -> Set a
Set.insert Cluster
c1 (Set Cluster -> Set Cluster)
-> (Set Cluster -> Set Cluster) -> Set Cluster -> Set Cluster
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Cluster -> Set Cluster -> Set Cluster
forall a. Ord a => a -> Set a -> Set a
Set.insert Cluster
c2  (Set Cluster -> Set Cluster) -> Set Cluster -> Set Cluster
forall a b. (a -> b) -> a -> b
$ Set Cluster
cs'
  where
    (Cluster
c, Set Cluster
cs') = Set Cluster -> (Cluster, Set Cluster)
forall a. Set a -> (a, Set a)
Set.deleteFindMax Set Cluster
cs
    (Cluster
c1, Cluster
c2) = Cluster -> (Cluster, Cluster)
subdivide Cluster
c

-- Keep splitting the initial cluster until there are 256 clusters, then return

-- a priority queue containing all 256.

clusters :: Int -> Image PixelRGB8 -> Set Cluster
clusters :: Int -> Image PixelRGB8 -> Set Cluster
clusters Int
maxCols Image PixelRGB8
img = Int -> Set Cluster
clusters' (Int
maxCols Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
  where
    clusters' :: Int -> Set Cluster
    clusters' :: Int -> Set Cluster
clusters' Int
0 = Cluster -> Set Cluster
forall a. a -> Set a
Set.singleton Cluster
c
    clusters' Int
n = Set Cluster -> Set Cluster
split (Int -> Set Cluster
clusters' (Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1))
    c :: Cluster
c = Image PixelRGB8 -> Cluster
initCluster Image PixelRGB8
img

-- Euclidean distance squared, between two pixels.

dist2Px :: PixelRGB8 -> PixelRGB8 -> Int
dist2Px :: PixelRGB8 -> PixelRGB8 -> Int
dist2Px (PixelRGB8 Pixel8
r1 Pixel8
g1 Pixel8
b1) (PixelRGB8 Pixel8
r2 Pixel8
g2 Pixel8
b2) = Int
drInt -> Int -> Int
forall a. Num a => a -> a -> a
*Int
dr Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
dgInt -> Int -> Int
forall a. Num a => a -> a -> a
*Int
dg Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
dbInt -> Int -> Int
forall a. Num a => a -> a -> a
*Int
db
  where
    (Int
dr, Int
dg, Int
db) =
      ( Pixel8 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Pixel8
r1 Int -> Int -> Int
forall a. Num a => a -> a -> a
- Pixel8 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Pixel8
r2
      , Pixel8 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Pixel8
g1 Int -> Int -> Int
forall a. Num a => a -> a -> a
- Pixel8 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Pixel8
g2
      , Pixel8 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Pixel8
b1 Int -> Int -> Int
forall a. Num a => a -> a -> a
- Pixel8 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Pixel8
b2 )

nearestColorIdx :: PixelRGB8 -> Vector PixelRGB8 -> Pixel8
nearestColorIdx :: PixelRGB8 -> Vector PixelRGB8 -> Pixel8
nearestColorIdx PixelRGB8
p Vector PixelRGB8
ps  = Int -> Pixel8
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> Pixel8) -> Int -> Pixel8
forall a b. (a -> b) -> a -> b
$ Vector Int -> Int
forall a. Ord a => Vector a -> Int
V.minIndex ((PixelRGB8 -> Int) -> Vector PixelRGB8 -> Vector Int
forall a b. (a -> b) -> Vector a -> Vector b
V.map (PixelRGB8 -> PixelRGB8 -> Int
`dist2Px` PixelRGB8
p) Vector PixelRGB8
ps)