{-# LANGUAGE ExistentialQuantification #-}
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE CPP #-}
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)
data PaletteCreationMethod =
MedianMeanCut
| Uniform
data PaletteOptions = PaletteOptions
{
PaletteOptions -> PaletteCreationMethod
paletteCreationMethod :: PaletteCreationMethod
, PaletteOptions -> Bool
enableImageDithering :: Bool
, PaletteOptions -> Int
paletteColorCount :: Int
}
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
}
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
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
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
Int
0
(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
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
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
uniformQuantization :: PaletteOptions -> Image PixelRGB8 -> (Image Pixel8, Palette)
uniformQuantization :: PaletteOptions
-> Image PixelRGB8 -> (Image Pixel8, Image PixelRGB8)
uniformQuantization PaletteOptions
opts Image PixelRGB8
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