Reduce memory usage of palette generator ⚡

palette-generator/Stylix/Main.hs

@@ -1,6 +1,6 @@
 import Ai.Evolutionary ( EvolutionConfig(EvolutionConfig), evolve )
-import Codec.Picture ( DynamicImage, Image(imageWidth, imageHeight), PixelRGB8(PixelRGB8), convertRGB8, pixelAt, readImage )
-import Data.Colour ( LAB, RGB(RGB), lab2rgb, rgb2lab )
+import Codec.Picture ( DynamicImage, Image, PixelRGB8, convertRGB8, readImage )
+import Data.Colour ( LAB, RGB(RGB), lab2rgb )
 import qualified Data.Vector as V
 import Stylix.Output ( makeOutputTable )
 import Stylix.Palette ( )
@@ -11,20 +11,11 @@ import Text.JSON ( encode )
 -- | Run the genetic algorithm to generate a palette from the given image.
 selectColours :: (Floating a, Real a)
               => String -- ^ Scheme type: "either", "light" or "dark"
-              -> V.Vector (LAB a) -- ^ Colours of the source image
+              -> Image PixelRGB8 -- ^ Source image
               -> IO (V.Vector (LAB a)) -- ^ Generated palette
 selectColours polarity image
   = evolve (polarity, image) (EvolutionConfig 1000 100 0.5 0.01)
 
--- | Convert a 'DynamicImage' to a simple 'V.Vector' of colours.
-unpackImage :: (Num a) => DynamicImage -> V.Vector (RGB a)
-unpackImage image = do
-  let image' = convertRGB8 image
-  x <- V.enumFromN 0 (imageWidth image')
-  y <- V.enumFromN 0 (imageHeight image')
-  let (PixelRGB8 r g b) = pixelAt image' x y
-  return $ RGB (fromIntegral r) (fromIntegral g) (fromIntegral b)
-
 -- | Load an image file.
 loadImage :: String -- ^ Path to the file
           -> IO DynamicImage
@@ -33,10 +24,9 @@ loadImage input = either error id <$> readImage input
 mainProcess :: (String, String, String) -> IO ()
 mainProcess (polarity, input, output) = do
   putStrLn $ "Processing " ++ input
-  image <- loadImage input
 
-  let labImage = V.map rgb2lab $ unpackImage image
-  palette <- selectColours polarity labImage
+  image <- loadImage input
+  palette <- selectColours polarity (convertRGB8 image)
   let outputTable = makeOutputTable $ V.map lab2rgb palette
 
   writeFile output $ encode outputTable

palette-generator/Stylix/Palette.hs

@@ -3,9 +3,10 @@
 module Stylix.Palette ( ) where
 
 import Ai.Evolutionary ( Species(..) )
+import Codec.Picture ( Image(imageWidth, imageHeight), PixelRGB8(PixelRGB8), pixelAt )
 import Data.Bifunctor ( second )
-import Data.Colour ( LAB(lightness), deltaE )
-import Data.Vector ( (!), (//) )
+import Data.Colour ( LAB(lightness), RGB(RGB), deltaE, rgb2lab )
+import Data.Vector ( (//) )
 import qualified Data.Vector as V
 import System.Random ( RandomGen, randomR )
 
@@ -25,16 +26,19 @@ taken enough colours for a new palette.
 alternatingZip :: V.Vector a -> V.Vector a -> V.Vector a
 alternatingZip = V.izipWith (\i a b -> if even i then a else b)
 
--- | Select a random item from a vector.
-randomFromVector :: (RandomGen r)
-                 => r -- ^ Random generator
-                 -> V.Vector a
-                 -> (a, r) -- ^ Chosen item, new random generator
-randomFromVector generator vector
-  = let (index, generator') = randomR (0, V.length vector - 1) generator
-     in (vector ! index, generator')
+-- | Select a random color from an image.
+randomFromImage :: (RandomGen r, Floating a, Num a, Ord a)
+                => r -- ^ Random generator
+                -> Image PixelRGB8
+                -> (LAB a, r) -- ^ Chosen color, new random generator
+randomFromImage generator image
+  = let (x, generator') = randomR (0, imageWidth image - 1) generator
+        (y, generator'') = randomR (0, imageHeight image - 1) generator'
+        (PixelRGB8 r g b) = pixelAt image x y
+        color = RGB (fromIntegral r) (fromIntegral g) (fromIntegral b)
+     in (rgb2lab color, generator'')
 
-instance (Floating a, Real a) => Species (String, (V.Vector (LAB a))) (V.Vector (LAB a)) where
+instance (Floating a, Real a) => Species (String, (Image PixelRGB8)) (V.Vector (LAB a)) where
   {- |
   Palettes in the initial population are created by randomly
   sampling 16 colours from the source image.
@@ -42,7 +46,7 @@ instance (Floating a, Real a) => Species (String, (V.Vector (LAB a))) (V.Vector
   generate (_, image) = generateColour 16
       where generateColour 0 generator = (generator, V.empty)
             generateColour n generator
-              = let (colour, generator') = randomFromVector generator image
+              = let (colour, generator') = randomFromImage generator image
                  in second (V.cons colour) $ generateColour (n - 1) generator'
 
   crossover _ generator a b = (generator, alternatingZip a b)
@@ -53,7 +57,7 @@ instance (Floating a, Real a) => Species (String, (V.Vector (LAB a))) (V.Vector
   -}
   mutate (_, image) generator palette
     = let (index, generator') = randomR (0, 15) generator
-          (colour, generator'') = randomFromVector generator' image
+          (colour, generator'') = randomFromImage generator' image
        in (generator'', palette // [(index, colour)])
 
   fitness (polarity, _) palette