{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE TypeFamilies     #-}

module PerceptualHash ( imgHash
                      , fileHash
                      , hammingDistance
                      ) where

import qualified Codec.Picture                 as JuicyPixels
import           Codec.Picture.WebP            (decodeRgb8)
import           Control.Monad.ST              (runST)
import           Data.Bits                     (Bits, popCount, shiftL, xor, (.|.))
import qualified Data.ByteString               as BS
import           Data.List                     (isSuffixOf)
import qualified Data.Vector.Generic           as V
import qualified Data.Vector.Storable          as VS
import           Data.Word                     (Word64, Word8)
import           Graphics.Image                (Array, Bilinear (..), Border (Edge, Reflect), Image,
                                                Pixel (PixelX, PixelY), RGB, RSU (..), VS, X, Y,
                                                convert, convolve, crop, makeImage, readImage,
                                                resize, transpose, (|*|))
import           Graphics.Image.Interface      (fromVector, toVector)
import qualified Graphics.Image.Interface      as Hip
import           Graphics.Image.Interface.Repa (fromRepaArrayS, toRepaArray)
import           Median                        (median)

-- | See
-- [wiki](https://en.wikipedia.org/wiki/Hamming_distance#Algorithm_example).
--
-- @since 0.1.4.0
{-# SPECIALIZE hammingDistance :: Word64 -> Word64 -> Int #-}
hammingDistance :: Bits a => a -> a -> Int
hammingDistance :: forall a. Bits a => a -> a -> Int
hammingDistance a
x a
y = a -> Int
forall a. Bits a => a -> Int
popCount (a
x a -> a -> a
forall a. Bits a => a -> a -> a
`xor` a
y)

dct32 :: (Floating e, Array arr Y e) => Image arr Y e
dct32 :: forall e arr. (Floating e, Array arr Y e) => Image arr Y e
dct32 = (Int, Int) -> ((Int, Int) -> Pixel Y e) -> Image arr Y e
forall arr cs e.
Array arr cs e =>
(Int, Int) -> ((Int, Int) -> Pixel cs e) -> Image arr cs e
makeImage (Int
32,Int
32) (Int, Int) -> Pixel Y e
gen
    where gen :: (Int, Int) -> Pixel Y e
gen (Int
i,Int
j) = e -> Pixel Y e
forall e. e -> Pixel Y e
PixelY (e -> Pixel Y e) -> e -> Pixel Y e
forall a b. (a -> b) -> a -> b
$ e -> e
forall a. Floating a => a -> a
sqrt(e
2e -> e -> e
forall a. Fractional a => a -> a -> a
/e
n) e -> e -> e
forall a. Num a => a -> a -> a
* e -> e
forall a. Floating a => a -> a
cos((Int -> e
forall a b. (Integral a, Num b) => a -> b
fromIntegral ((Int
2Int -> Int -> Int
forall a. Num a => a -> a -> a
*Int
i) Int -> Int -> Int
forall a. Num a => a -> a -> a
* (Int
jInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1)) e -> e -> e
forall a. Num a => a -> a -> a
* e
forall a. Floating a => a
pi)e -> e -> e
forall a. Fractional a => a -> a -> a
/(e
2e -> e -> e
forall a. Num a => a -> a -> a
*e
n))
          n :: e
n = e
32

idMat :: (Fractional e, Array arr X e) => Image arr X e
idMat :: forall e arr. (Fractional e, Array arr X e) => Image arr X e
idMat = (Int, Int) -> ((Int, Int) -> Pixel X e) -> Image arr X e
forall arr cs e.
Array arr cs e =>
(Int, Int) -> ((Int, Int) -> Pixel cs e) -> Image arr cs e
makeImage (Int
7,Int
7) (\(Int, Int)
_ -> e -> Pixel X e
forall e. e -> Pixel X e
PixelX (e
1e -> e -> e
forall a. Fractional a => a -> a -> a
/e
49))

{-# INLINE meanFilter #-}
meanFilter :: (Fractional e, Array arr X e, Array arr cs e) => Image arr cs e -> Image arr cs e
meanFilter :: forall e arr cs.
(Fractional e, Array arr X e, Array arr cs e) =>
Image arr cs e -> Image arr cs e
meanFilter = {-# SCC "meanFilter" #-} Border (Pixel cs e)
-> Image arr X e -> Image arr cs e -> Image arr cs e
forall arr e cs.
(Array arr X e, Array arr cs e) =>
Border (Pixel cs e)
-> Image arr X e -> Image arr cs e -> Image arr cs e
convolve Border (Pixel cs e)
forall px. Border px
Reflect Image arr X e
forall e arr. (Fractional e, Array arr X e) => Image arr X e
idMat

{-# INLINE size32 #-}
size32 :: Array arr cs e => Image arr cs e -> Image arr cs e
size32 :: forall arr cs e. Array arr cs e => Image arr cs e -> Image arr cs e
size32 = Bilinear
-> Border (Pixel cs e)
-> (Int, Int)
-> Image arr cs e
-> Image arr cs e
forall method arr cs e.
(Interpolation method, Array arr cs e) =>
method
-> Border (Pixel cs e)
-> (Int, Int)
-> Image arr cs e
-> Image arr cs e
resize Bilinear
Bilinear Border (Pixel cs e)
forall px. Border px
Edge (Int
32,Int
32)

crop8 :: Array arr cs e => Image arr cs e -> Image arr cs e
crop8 :: forall arr cs e. Array arr cs e => Image arr cs e -> Image arr cs e
crop8 = (Int, Int) -> (Int, Int) -> Image arr cs e -> Image arr cs e
forall arr cs e.
Array arr cs e =>
(Int, Int) -> (Int, Int) -> Image arr cs e -> Image arr cs e
crop (Int
0,Int
0) (Int
8,Int
8)

medianImmut :: (Ord e, Fractional e, V.Vector v e) => v e -> e
medianImmut :: forall e (v :: * -> *).
(Ord e, Fractional e, Vector v e) =>
v e -> e
medianImmut v e
v = (forall s. ST s e) -> e
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s e) -> e) -> (forall s. ST s e) -> e
forall a b. (a -> b) -> a -> b
$
    Mutable v s e -> ST s e
forall (m :: * -> *) (v :: * -> * -> *) e.
(PrimMonad m, MVector v e, Ord e, Fractional e) =>
v (PrimState m) e -> m e
median (Mutable v s e -> ST s e) -> ST s (Mutable v s e) -> ST s e
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< v e -> ST s (Mutable v (PrimState (ST s)) e)
forall (m :: * -> *) (v :: * -> *) a.
(PrimMonad m, Vector v a) =>
v a -> m (Mutable v (PrimState m) a)
V.thaw v e
v

dct :: (Floating e, Array arr Y e) => Image arr Y e -> Image arr Y e
dct :: forall e arr.
(Floating e, Array arr Y e) =>
Image arr Y e -> Image arr Y e
dct Image arr Y e
img = Image arr Y e
forall e arr. (Floating e, Array arr Y e) => Image arr Y e
dct32 Image arr Y e -> Image arr Y e -> Image arr Y e
forall arr cs e.
Array arr cs e =>
Image arr cs e -> Image arr cs e -> Image arr cs e
|*| Image arr Y e
img Image arr Y e -> Image arr Y e -> Image arr Y e
forall arr cs e.
Array arr cs e =>
Image arr cs e -> Image arr cs e -> Image arr cs e
|*| Image arr Y e -> Image arr Y e
forall arr cs e. Array arr cs e => Image arr cs e -> Image arr cs e
transpose Image arr Y e
forall e arr. (Floating e, Array arr Y e) => Image arr Y e
dct32

{-# INLINE imgHash #-}
-- | DCT based hash. See
-- [Zauner](https://www.phash.org/docs/pubs/thesis_zauner.pdf).
--
-- It is suggested that you use this with the Repa backend.
imgHash :: (Ord e, Floating e, Array arr Y e, Array arr X e, V.Vector (Hip.Vector arr) Bool, V.Vector (Hip.Vector arr) e) => Image arr Y e -> Word64
imgHash :: forall e arr.
(Ord e, Floating e, Array arr Y e, Array arr X e,
 Vector (Vector arr) Bool, Vector (Vector arr) e) =>
Image arr Y e -> Word64
imgHash = Vector arr Bool -> Word64
forall (v :: * -> *). Vector v Bool => v Bool -> Word64
asWord64 (Vector arr Bool -> Word64)
-> (Image arr Y e -> Vector arr Bool) -> Image arr Y e -> Word64
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Vector arr e -> Vector arr Bool
forall e (v :: * -> *).
(Fractional e, Vector v e, Vector v Bool, Ord e) =>
v e -> v Bool
aboveMed (Vector arr e -> Vector arr Bool)
-> (Image arr Y e -> Vector arr e)
-> Image arr Y e
-> Vector arr Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Pixel Y e -> e) -> Vector arr (Pixel Y e) -> Vector arr e
forall (v :: * -> *) a b.
(Vector v a, Vector v b) =>
(a -> b) -> v a -> v b
V.map (\(PixelY e
x) -> e
x) (Vector arr (Pixel Y e) -> Vector arr e)
-> (Image arr Y e -> Vector arr (Pixel Y e))
-> Image arr Y e
-> Vector arr e
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Image arr Y e -> Vector arr (Pixel Y e)
forall arr cs e.
Array arr cs e =>
Image arr cs e -> Vector arr (Pixel cs e)
toVector (Image arr Y e -> Vector arr (Pixel Y e))
-> (Image arr Y e -> Image arr Y e)
-> Image arr Y e
-> Vector arr (Pixel Y e)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Image arr Y e -> Image arr Y e
forall arr cs e. Array arr cs e => Image arr cs e -> Image arr cs e
crop8 (Image arr Y e -> Image arr Y e)
-> (Image arr Y e -> Image arr Y e)
-> Image arr Y e
-> Image arr Y e
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Image arr Y e -> Image arr Y e
forall e arr.
(Floating e, Array arr Y e) =>
Image arr Y e -> Image arr Y e
dct (Image arr Y e -> Image arr Y e)
-> (Image arr Y e -> Image arr Y e)
-> Image arr Y e
-> Image arr Y e
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Image arr Y e -> Image arr Y e
forall arr cs e. Array arr cs e => Image arr cs e -> Image arr cs e
size32 (Image arr Y e -> Image arr Y e)
-> (Image arr Y e -> Image arr Y e)
-> Image arr Y e
-> Image arr Y e
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Image arr Y e -> Image arr Y e
forall e arr cs.
(Fractional e, Array arr X e, Array arr cs e) =>
Image arr cs e -> Image arr cs e
meanFilter

asWord64 :: V.Vector v Bool => v Bool -> Word64
asWord64 :: forall (v :: * -> *). Vector v Bool => v Bool -> Word64
asWord64 = (Word64 -> Bool -> Word64) -> Word64 -> v Bool -> Word64
forall (v :: * -> *) b a.
Vector v b =>
(a -> b -> a) -> a -> v b -> a
V.foldl' (\Word64
acc Bool
x -> (Word64
acc Word64 -> Int -> Word64
forall a. Bits a => a -> Int -> a
`shiftL` Int
1) Word64 -> Word64 -> Word64
forall a. Bits a => a -> a -> a
.|. Bool -> Word64
boolToWord64 Bool
x) Word64
0
    where boolToWord64 :: Bool -> Word64
          boolToWord64 :: Bool -> Word64
boolToWord64 Bool
False = Word64
0
          boolToWord64 Bool
True  = Word64
1

aboveMed :: (Fractional e, V.Vector v e, V.Vector v Bool, Ord e) => v e -> v Bool
aboveMed :: forall e (v :: * -> *).
(Fractional e, Vector v e, Vector v Bool, Ord e) =>
v e -> v Bool
aboveMed v e
v =
    let med :: e
med = v e -> e
forall e (v :: * -> *).
(Ord e, Fractional e, Vector v e) =>
v e -> e
medianImmut v e
v
    in (e -> Bool) -> v e -> v Bool
forall (v :: * -> *) a b.
(Vector v a, Vector v b) =>
(a -> b) -> v a -> v b
V.map (e -> e -> Bool
forall a. Ord a => a -> a -> Bool
<e
med) v e
v

fileWebp :: FilePath -> IO (Image VS RGB Word8)
fileWebp :: String -> IO (Image VS RGB Word8)
fileWebp String
fp = do
    ByteString
contents <- String -> IO ByteString
BS.readFile String
fp
    let (JuicyPixels.Image Int
m Int
n Vector (PixelBaseComponent PixelRGB8)
pixels) = ByteString -> Image PixelRGB8
decodeRgb8 ByteString
contents
    Image VS RGB Word8 -> IO (Image VS RGB Word8)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Image VS RGB Word8 -> IO (Image VS RGB Word8))
-> Image VS RGB Word8 -> IO (Image VS RGB Word8)
forall a b. (a -> b) -> a -> b
$ (Int, Int) -> Vector VS (Pixel RGB Word8) -> Image VS RGB Word8
forall arr cs e.
Array arr cs e =>
(Int, Int) -> Vector arr (Pixel cs e) -> Image arr cs e
fromVector (Int
m, Int
n) (Vector VS (Pixel RGB Word8) -> Image VS RGB Word8)
-> Vector VS (Pixel RGB Word8) -> Image VS RGB Word8
forall a b. (a -> b) -> a -> b
$ Vector Word8 -> Vector (Pixel RGB Word8)
forall a b. (Storable a, Storable b) => Vector a -> Vector b
VS.unsafeCast Vector Word8
Vector (PixelBaseComponent PixelRGB8)
pixels

readWebp :: FilePath -> IO (Image VS Y Double)
readWebp :: String -> IO (Image VS Y Double)
readWebp = (Image VS RGB Word8 -> Image VS Y Double)
-> IO (Image VS RGB Word8) -> IO (Image VS Y Double)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Image VS RGB Word8 -> Image VS Y Double
forall cs e cs' e' arr.
(Convertible cs e, ToYA cs' e', ToRGBA cs' e', Array arr cs' e',
 Array arr cs e) =>
Image arr cs' e' -> Image arr cs e
convert (IO (Image VS RGB Word8) -> IO (Image VS Y Double))
-> (String -> IO (Image VS RGB Word8))
-> String
-> IO (Image VS Y Double)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> IO (Image VS RGB Word8)
fileWebp

-- | @since 0.1.5.0
fileHashWebp :: FilePath -> IO Word64
fileHashWebp :: String -> IO Word64
fileHashWebp = (Image VS Y Double -> Word64)
-> IO (Image VS Y Double) -> IO Word64
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Image RSU Y Double -> Word64
forall e arr.
(Ord e, Floating e, Array arr Y e, Array arr X e,
 Vector (Vector arr) Bool, Vector (Vector arr) e) =>
Image arr Y e -> Word64
imgHash (Image RSU Y Double -> Word64)
-> (Image VS Y Double -> Image RSU Y Double)
-> Image VS Y Double
-> Word64
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Image VS Y Double -> Image RSU Y Double
convRepa) (IO (Image VS Y Double) -> IO Word64)
-> (String -> IO (Image VS Y Double)) -> String -> IO Word64
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> IO (Image VS Y Double)
readWebp
    -- faster
    where convRepa :: Image VS Y Double -> Image RSU Y Double
convRepa = Array U DIM2 (Pixel Y Double) -> Image RSU Y Double
forall r cs e.
Source r (Pixel cs e) =>
Array r DIM2 (Pixel cs e) -> Image RSU cs e
fromRepaArrayS (Array U DIM2 (Pixel Y Double) -> Image RSU Y Double)
-> (Image VS Y Double -> Array U DIM2 (Pixel Y Double))
-> Image VS Y Double
-> Image RSU Y Double
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Image VS Y Double -> Array U DIM2 (Pixel Y Double)
forall arr cs e.
Array arr cs e =>
Image arr cs e -> Array U DIM2 (Pixel cs e)
toRepaArray

fileHash :: FilePath -> IO (Either String Word64)
fileHash :: String -> IO (Either String Word64)
fileHash String
fp | String
".webp" String -> String -> Bool
forall a. Eq a => [a] -> [a] -> Bool
`isSuffixOf` String
fp = Word64 -> Either String Word64
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Word64 -> Either String Word64)
-> IO Word64 -> IO (Either String Word64)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> String -> IO Word64
fileHashWebp String
fp
            | Bool
otherwise = String -> IO (Either String Word64)
fileHashHip String
fp

fileHashHip :: FilePath -> IO (Either String Word64)
fileHashHip :: String -> IO (Either String Word64)
fileHashHip = (Either String (Image RSU Y Double) -> Either String Word64)
-> IO (Either String (Image RSU Y Double))
-> IO (Either String Word64)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((Image RSU Y Double -> Word64)
-> Either String (Image RSU Y Double) -> Either String Word64
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Image RSU Y Double -> Word64
forall e arr.
(Ord e, Floating e, Array arr Y e, Array arr X e,
 Vector (Vector arr) Bool, Vector (Vector arr) e) =>
Image arr Y e -> Word64
imgHash :: Image RSU Y Double -> Word64)) (IO (Either String (Image RSU Y Double))
 -> IO (Either String Word64))
-> (String -> IO (Either String (Image RSU Y Double)))
-> String
-> IO (Either String Word64)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> IO (Either String (Image RSU Y Double))
forall arr cs e.
(Array VS cs e, Array arr cs e,
 Readable (Image VS cs e) InputFormat) =>
String -> IO (Either String (Image arr cs e))
readImage