Voice.hs

{-# LANGUAGE ApplicativeDo     #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RecordWildCards   #-}
{-|
  Reanimate can automatically synchronize animations to your voice if you have
  a transcript and an audio recording. This works with the help of Gentle
  (<https://lowerquality.com/gentle/>). Accuracy is not perfect but it is pretty
  close, and it is by far the easiest way of adding narration to an animation.
-}
module Reanimate.Voice
  ( Transcript(..)
  , TWord(..)
  , Phone(..)
  , findWord                -- :: Transcript -> [Text] -> Text -> TWord
  , findWords               -- :: Transcript -> [Text] -> Text -> [TWord]
  , loadTranscript          -- :: FilePath -> Transcript
  , fakeTranscript          -- :: Text -> Transcript
  , splitTranscript         -- :: Transcript -> SVG -> [(SVG, TWord)]
  , annotateWithTranscript  -- :: Transcript -> Scene s ()
  )
where

import           Control.Monad       (forM_)
import           Data.Aeson
import           Data.Char           (isAlphaNum, isSpace)
import           Data.List           (sortOn)
import           Data.Map            (Map)
import qualified Data.Map            as Map
import           Data.Maybe          (listToMaybe)
import           Data.Text           (Text)
import qualified Data.Text           as T
import qualified Data.Text.IO        as T
import           Reanimate.Animation (SVG, staticFrame)
import           Reanimate.Constants (defaultStrokeWidth, screenHeight, screenWidth)
import           Reanimate.LaTeX     (latex, latexChunks)
import           Reanimate.Misc      (withTempFile)
import           Reanimate.Scene     (Scene, play, waitUntil)
import           Reanimate.Svg       (mkGroup, scale, translate, withStrokeColor, withStrokeWidth)
import           System.Directory    (doesFileExist)
import           System.Exit         (ExitCode (ExitFailure, ExitSuccess))
import           System.FilePath     (replaceExtension)
import           System.IO.Unsafe    (unsafePerformIO)
import           System.Process      (rawSystem, showCommandForUser)

-- | Aligned transcript. Contains the transcript text as well as
--   timing data for each word.
data Transcript = Transcript
  { transcriptText  :: Text
  , transcriptKeys  :: Map Text Int
  , transcriptWords :: [TWord]
  } deriving (Show)

instance FromJSON Transcript where
  parseJSON = withObject "transcript" $ \o ->
    Transcript <$> o .: "transcript" <*> pure Map.empty <*> o .: "words"

-- | Spoken word. Includes information about when it was spoken,
--   its duration, and its phonemes.
data TWord = TWord
  { wordAligned     :: Text
  , wordCase        :: Text
  , wordStart       :: Double -- ^ Start of pronunciation in seconds
  , wordStartOffset :: Int    -- ^ Character index of word in transcript
  , wordEnd         :: Double -- ^ End of pronunciation in seconds
  , wordEndOffset   :: Int    -- ^ Last character index of word in transcript
  , wordPhones      :: [Phone]
  , wordReference   :: Text   -- ^ The word being pronounced.
  } deriving (Show)

instance FromJSON TWord where
  parseJSON = withObject "word" $ \o ->
    TWord
      <$> o
      .:? "alignedWord"
      .!= T.empty
      <*> o
      .:  "case"
      <*> o
      .:? "start"
      .!= 0
      <*> o
      .:  "startOffset"
      <*> o
      .:? "end"
      .!= 0
      <*> o
      .:  "endOffset"
      <*> o
      .:? "phones"
      .!= []
      <*> o
      .:  "word"

-- | Phoneme type
data Phone = Phone
  { phoneDuration :: Double
  , phoneType     :: Text
  } deriving (Show)

instance FromJSON Phone where
  parseJSON =
    withObject "phone" $ \o -> Phone <$> o .: "duration" <*> o .: "phone"

-- | Locate the first word that occurs after all the given keys.
--   An error is thrown if no such word exists. An error is thrown
--   if the keys do not exist in the transcript.
findWord :: Transcript -> [Text] -> Text -> TWord
findWord t keys w = case listToMaybe (findWords t keys w) of
  Nothing    -> error $ "Word not in transcript: " ++ show (keys, w)
  Just tword -> tword

-- | Locate all words that occur after all the given keys.
--   May return an empty list. An error is thrown
--   if the keys do not exist in the transcript.
findWords :: Transcript -> [Text] -> Text -> [TWord]
findWords t [] wd =
  [ tword | tword <- transcriptWords t, wordReference tword == wd ]
findWords t (key : keys) wd =
  [ tword
  | tword <- findWords t keys wd
  , wordStartOffset tword >= Map.findWithDefault badKey key (transcriptKeys t)
  ]
  where badKey = error $ "Missing transcript key: " ++ show key

-- | Loading a transcript does three things depending on which files are available
--   with the same basename as the input argument:
--   1. If a JSON file is available, it is parsed and returned.
--   2. If an audio file is available, reanimate tries to align it by calling out to
--      Gentle on localhost:8765/. If Gentle is not running, an error will be thrown.
--   3. If only the text transcript is available, a fake transcript is returned,
--      with timings roughly at 120 words per minute.
loadTranscript :: FilePath -> Transcript
loadTranscript path = unsafePerformIO $ do
  rawTranscript <- T.readFile path
  let keys           = parseTranscriptKeys rawTranscript
      trimTranscript = cutoutKeys keys rawTranscript
  hasJSON    <- doesFileExist jsonPath
  transcript <- if hasJSON
    then do
      mbT <- decodeFileStrict jsonPath
      case mbT of
        Nothing -> error "bad json"
        Just t  -> pure t
    else do
      hasAudio <- findWithExtension path audioExtensions
      case hasAudio of
        Nothing        -> return $ fakeTranscript' trimTranscript
        Just audioPath -> withTempFile "txt" $ \txtPath -> do
          T.writeFile txtPath trimTranscript
          runGentleForcedAligner audioPath txtPath
          mbT <- decodeFileStrict jsonPath
          case mbT of
            Nothing -> error "bad json"
            Just t  -> pure t
  pure $ transcript { transcriptKeys = finalizeKeys keys }
 where
  jsonPath        = replaceExtension path "json"
  audioExtensions = ["mp3", "m4a", "flac"]

parseTranscriptKeys :: Text -> Map Text Int
parseTranscriptKeys = worker Map.empty 0
 where
  worker keys offset txt = case T.uncons txt of
    Nothing -> keys
    Just ('[', cs) ->
      let key       = T.takeWhile (/= ']') cs
          newOffset = T.length key + 2
      in  worker (Map.insert key offset keys)
                 (offset + newOffset)
                 (T.drop newOffset txt)
    Just (_, cs) -> worker keys (offset + 1) cs

finalizeKeys :: Map Text Int -> Map Text Int
finalizeKeys = Map.fromList . worker 0 . sortOn snd . Map.toList
 where
  worker _offset [] = []
  worker offset ((key, at) : rest) =
    (key, at - offset) : worker (offset + T.length key + 2) rest

cutoutKeys :: Map Text Int -> Text -> Text
cutoutKeys keys = T.concat . worker 0 (sortOn snd (Map.toList keys))
 where
  worker _offset [] txt = [txt]
  worker offset ((key, at) : xs) txt =
    let keyLen          = T.length key + 2
        (before, after) = T.splitAt (at - offset) txt
    in  before : worker (at + keyLen) xs (T.drop keyLen after)

findWithExtension :: FilePath -> [String] -> IO (Maybe FilePath)
findWithExtension _path []       = return Nothing
findWithExtension path  (e : es) = do
  let newPath = replaceExtension path e
  hasFile <- doesFileExist newPath
  if hasFile then return (Just newPath) else findWithExtension path es

runGentleForcedAligner :: FilePath -> FilePath -> IO ()
runGentleForcedAligner audioFile transcriptFile = do
  ret <- rawSystem prog args
  case ret of
    ExitSuccess -> return ()
    ExitFailure e ->
      error
        $  "Gentle forced aligner failed with: "
        ++ show e
        ++ "\nIs it running locally on port 8765?"
        ++ "\nCommand: "
        ++ showCommandForUser prog args
 where
  prog = "curl"
  args =
    [ "--silent"
    , "--form"
    , "audio=@" ++ audioFile
    , "--form"
    , "transcript=@" ++ transcriptFile
    , "--output"
    , replaceExtension audioFile "json"
    , "http://localhost:8765/transcriptions?async=false"
    ]

data Token = TokenWord Int Int Text | TokenComma | TokenPeriod | TokenParagraph
  deriving (Show)

lexText :: Text -> [Token]
lexText = worker 0
 where
  worker offset txt = case T.uncons txt of
    Nothing -> []
    Just (c, cs)
      | isSpace c
      -> let (w, rest) = T.span (== '\n') txt
         in  if T.length w >= 3
               then TokenParagraph : worker (offset + T.length w) rest
               else worker (offset + 1) cs
      | c == '.'
      -> TokenPeriod : worker (offset + 1) cs
      | c == ','
      -> TokenComma : worker (offset + 1) cs
      | isWord c
      -> let (w, rest) = T.span isWord txt
             newOffset = offset + T.length w
         in  TokenWord offset newOffset w : worker newOffset rest
      | otherwise
      -> worker (offset + 1) cs
  isWord c = isAlphaNum c || c `elem` ['\'', '-']

-- | Fake transcript timings at roughly 120 words per minute.
fakeTranscript :: Text -> Transcript
fakeTranscript rawTranscript =
  let keys = parseTranscriptKeys rawTranscript
      t    = fakeTranscript' (cutoutKeys keys rawTranscript)
  in  t { transcriptKeys = finalizeKeys keys }

fakeTranscript' :: Text -> Transcript
fakeTranscript' input = Transcript { transcriptText  = input
                                   , transcriptKeys  = Map.empty
                                   , transcriptWords = worker 0 (lexText input)
                                   }
 where
  worker _now []             = []
  worker now  (token : rest) = case token of
    TokenWord start end w ->
      let dur = realToFrac (end - start) * 0.1
      in  TWord { wordAligned     = T.toLower w
                , wordCase        = "success"
                , wordStart       = now
                , wordStartOffset = start
                , wordEnd         = now + dur
                , wordEndOffset   = end
                , wordPhones      = []
                , wordReference   = w
                }
            : worker (now + dur) rest
    TokenComma     -> worker (now + commaPause) rest
    TokenPeriod    -> worker (now + periodPause) rest
    TokenParagraph -> worker (now + paragraphPause) rest
  paragraphPause = 0.5
  commaPause     = 0.1
  periodPause    = 0.2

-- | Convert the transcript text to an SVG image using LaTeX and associate
--   each word image with its timing information.
splitTranscript :: Transcript -> [(SVG, TWord)]
splitTranscript Transcript {..} =
  [ (svg, tword)
  | tword@TWord {..} <- transcriptWords
  , let wordLength  = wordEndOffset - wordStartOffset
        [_, svg, _] = latexChunks
          [ T.take wordStartOffset transcriptText
          , T.take wordLength (T.drop wordStartOffset transcriptText)
          , T.drop wordEndOffset transcriptText
          ]
  ]

-- | Helper function for rendering a transcript.
annotateWithTranscript :: Transcript -> Scene s ()
annotateWithTranscript t = forM_ (transcriptWords t) $ \tword -> do
  let svg = scale 1 $ latex (wordReference tword)
  waitUntil (wordStart tword)
  let dur = wordEnd tword - wordStart tword
  play $ staticFrame dur $ position $ outline svg
 where
  position = translate (-screenWidth / 2) (-screenHeight / 2)
  outline txt = mkGroup
    [ withStrokeWidth (defaultStrokeWidth * 10) $ withStrokeColor "white" txt
    , withStrokeWidth 0 txt
    ]