CODEWITHSUNDEEP
haskelloptimization
Yofe
Yofe

Reputation: 447

How can I optimize this haskell function

I need to find the closest color in a palette ps to a given color p. How do I make the function nearestColor as fast as possible, without changing the type of Pixel8 or PixelRGB8. So far I have tried inlining.

import qualified Data.Vector as V

type Pixel8 = Word8    

data PixelRGB8 = PixelRGB8 {-# UNPACK #-} !Pixel8 -- Red
                           {-# UNPACK #-} !Pixel8 -- Green
                           {-# UNPACK #-} !Pixel8 -- Blue
           deriving (Eq, Ord, Show)

nearestColor :: PixelRGB8 -> Vector PixelRGB8 -> PixelRGB8
nearestColor p ps = snd $ V.minimumBy comp ds
  where
    ds = V.map (\px -> (dist2Px px p, px)) ps
    comp a b = fst a `compare` fst b

dist2Px :: PixelRGB8 -> PixelRGB8 -> Int
dist2Px (PixelRGB8 r1 g1 b1) (PixelRGB8 r2 g2 b2) = dr*dr + dg*dg + db*db
  where
    (dr, dg, db) =
      ( fromIntegral r1 - fromIntegral r2
      , fromIntegral g1 - fromIntegral g2
      , fromIntegral b1 - fromIntegral b2 )

Upvotes: 3

Views: 161

Answers (1)

leftaroundabout
leftaroundabout

Reputation: 120711

If you want to use a single palette and request different colours, I'd first flip your signature:

type Palette = V.Vector PixelRGB8
nearestColor :: Palette -> PixelRGB8 -> PixelRGB8

That facilitates partial application, and allows the palette configuration to be memoised.

Next, you want to do that: re-store the palette in a data structure suitable for fast lookup. Since you're basically interested in Euclidean distance in ℝ3 (BTW not really ideal for colour comparison), this is a very common problem. A classic structure is the k-d tree, which has long been used for such a nearest-neighbour search. There's sure enough a Haskell library available, which quite a convenient interface for you:

import qualified Data.Trees.KdTree a s KD

instance KD.Point PixelRGB where
  dimension _ = 3
  coord 0 (PixelRGB r _ _) = fromIntegral r
  coord 1 (PixelRGB _ g _) = fromIntegral g
  coord 2 (PixelRGB _ _ b) = fromIntegral b
  dist2 = fromIntegral . dist2Px

Then we can transform a palette into such a tree:

type FastPalette = KD.KdTree PixelRGB8
accelPalette :: Palette -> FastPalette
accelPalette = KD.fromList . V.toList

And finally just use the library-provided next neighbour search:

nearestColor palette = fromJust . KD.nearestNeighbor fpal
 where fpal = accelPalette palette

Upvotes: 6

Related Questions