packages feed

HarmTrace-1.0: src/Main.hs

{-# OPTIONS_GHC -Wall #-}
{-# LANGUAGE FlexibleInstances                #-}

module Main where

-- Libs
import System.Console.ParseArgs hiding (args) 

-- Music stuff
import HarmTrace.HarmTrace
import HarmTrace.IO.Main
import HarmTrace.IO.Errors
import HarmTrace.IO.PrintTree
import HarmTrace.HAnTree.ToHAnTree (gTreeHead)
import HarmTrace.Matching.Standard
import HarmTrace.Matching.GuptaNishimura (getLCES)
import HarmTrace.Matching.Alignment (getAlignDist, alignChordLab
                                    , pPrintV, alignHAnChord)
import HarmTrace.Audio.Harmonize    ( simpleAnnotator, headAnnotator
                                    , harmonyAnnotator)
import HarmTrace.Audio.ChordTypes (AudioFeat, ChordAnnotation, TimedData)
import HarmTrace.Base.MusicRep (Key)

import Data.List (delete)
import System.FilePath (takeFileName)

--------------------------------------------------------------------------------
-- Command-line arguments
--------------------------------------------------------------------------------

data MyArgs = SourceInputString | SourceInputFile | TargetInputFile 
            | InputDir | OpMode | Print | MaxErrorRate | BinaryOut | BinaryIn
            | PrintIns | Grammar | TargetKeyInputFile
            | SourceKeyInputFile | AnnotationKeyInputDir| GroundTruthInputFile 
            | GroundTruthInputDir
            deriving (Eq, Ord, Show)

myArgs :: [Arg MyArgs]
myArgs = [
          Arg { argIndex = MaxErrorRate,
                argAbbr  = Just 'e',
                argName  = Just "max-err",
                argData  = argDataOptional "float" ArgtypeFloat,
                argDesc  = "Ignore pieces with higher error rate for diff"
              },
          Arg { argIndex = SourceInputString,
                argAbbr  = Just 'c',
                argName  = Just "chords",
                argData  = argDataOptional "string" ArgtypeString,
                argDesc  = "Input Chord Sequence to parse"
              },
          Arg { argIndex = SourceInputFile,
                argAbbr  = Just '1',
                argName  = Just "sfile",
                argData  = argDataOptional "filepath" ArgtypeString,
                argDesc  = "Input file (source for diff)"
              },
          Arg { argIndex = TargetInputFile,
                argAbbr  = Just '2',
                argName  = Just "tfile",
                argData  = argDataOptional "filepath" ArgtypeString,
                argDesc  = "Input file (target for diff)"
              },
          Arg { argIndex = SourceKeyInputFile,
                argAbbr  = Just '3',
                argName  = Just "skey",
                argData  = argDataOptional "filepath" ArgtypeString,
                argDesc  = "Source input ground-truth key annotation file"
              },
          Arg { argIndex = TargetKeyInputFile,
                argAbbr  = Just '4',
                argName  = Just "tkey",
                argData  = argDataOptional "filepath" ArgtypeString,
                argDesc  = "Target input ground-truth key annotation file"
              },
          Arg { argIndex = GroundTruthInputFile,
                argAbbr  = Just '5',
                argName  = Just "gt",
                argData  = argDataOptional "filepath" ArgtypeString,
                argDesc  = "Source input ground-truth chord annotation file"
              },
          Arg { argIndex = GroundTruthInputDir,
                argAbbr  = Just 'a',
                argName  = Just "gt-dir",
                argData  = argDataOptional "dir" ArgtypeString,
                argDesc  = "Ground-truth Annotation directory"
              },
          Arg { argIndex = InputDir,
                argAbbr  = Just 'd',
                argName  = Just "dir",
                argData  = argDataOptional "dir" ArgtypeString,
                argDesc  = "Input directory (process all files within)"
              },
          Arg { argIndex = AnnotationKeyInputDir,
                argAbbr  = Just 'k',
                argName  = Just "key-dir",
                argData  = argDataOptional "dir" ArgtypeString,
                argDesc  = "Ground-truth Key annotation directory"
              },              
          Arg { argIndex = OpMode,
                argAbbr  = Just 'm',
                argName  = Just "mode",
                argData  = argDataRequired "string" ArgtypeString,
                argDesc  = 
                  "Mode: parse|stdiff|lces-s|lcessim|hanlign|\n" ++ 
  "                                 align|harm-cr|head-cr|simple-cr"
              },
          Arg { argIndex = Grammar,
                argAbbr  = Just 'g',
                argName  = Just "grammar",
                argData  = argDataRequired "string" ArgtypeString,
                argDesc  = "Grammar to use (jazz|pop)"
              },
          Arg { argIndex = Print,
                argAbbr  = Just 'p',
                argName  = Just "print",
                argData  = Nothing,
                argDesc  = "Set this flag to print a .png of the parse"
              },
          Arg { argIndex = PrintIns,
                argAbbr  = Just 's',
                argName  = Just "print-insertions",
                argData  = Nothing,
                argDesc  = "Set this flag to show inserted nodes"
              },
          Arg { argIndex = BinaryOut,
                argAbbr  = Just 'o',
                argName  = Just "out",
                argData  = argDataOptional "filepath" ArgtypeString,
                argDesc  = "Output binary file for parsing results"
              },
          Arg { argIndex = BinaryIn,
                argAbbr  = Just 'i',
                argName  = Just "in",
                argData  = argDataOptional "filepath" ArgtypeString,
                argDesc  = "Input binary file for matching"
              }
         ]

--------------------------------------------------------------------------------
-- Main
--------------------------------------------------------------------------------

data SourceData = Audio | ChordLab | Annotation
  
getDataType :: Args MyArgs -> SourceData
getDataType args
  | gotArg args SourceKeyInputFile     = Annotation -- for single files
  | gotArg args AnnotationKeyInputDir  = Annotation -- for directory reading
  | otherwise                          = ChordLab   -- add audio support

-- by default all post processing operations are executed             
defaultOpts :: [PPOption]
defaultOpts  = [ RemovePDPT    , RemoveInsertions
               , MergeDelChords, ExpandChordDurations ]    
  
err1, err2, err3, err4 :: String
err1 = "Use a source file, or a directory."
err2 = "Use a source file and a target file, or a directory."
err3 = "Use a source file and optionally a target file."
err4 = "Use an audio-feature location and a ground-truth file, "++
       "or an audio-feature directory and a ground-truth directory."

main :: IO ()
main = do args <- parseArgsIO ArgsComplete myArgs
          let mode = getRequiredArg args OpMode
              grmS = getRequiredArg args Grammar
              prnt = gotArg args Print
              opts = if gotArg args PrintIns 
                     then delete RemoveInsertions defaultOpts else defaultOpts
              gram = case grmS of
                       "jazz" -> GrammarEx Jazz
                       "pop"  -> GrammarEx Pop
                       s      -> usageError args ("Unknown grammar: " ++ s)
          case mode of
            "parse"     -> mainParse args opts prnt           gram
            "stdiff"    -> mainMatch args opts False STDiff   gram
            "hanlign"   -> mainMatch args opts prnt  HAnAlign gram
            "lces-s"    -> mainMatch args opts prnt  LCESsize gram
            "lcessim"   -> mainMatch args opts prnt  LCESsim  gram
            "align"     -> mainMatch args opts prnt  Align    gram
            "head-cr"   -> mainChordRec headAnnotator   args prnt
            "simple-cr" -> mainChordRec simpleAnnotator args prnt
            "harm-cr"   -> mainChordRec (harmonyAnnotator gram) args prnt
            s           -> usageError args ("Unknown mode: " ++ s)


mainParse :: Args MyArgs -> [PPOption] -> Bool -> GrammarEx -> IO ()
mainParse args o p (GrammarEx g) =
  do let cStr = getArgString args SourceInputString
         mf1  = getArgString args SourceInputFile
         ky   = getArgString args SourceKeyInputFile
         bOut = getArgString args BinaryOut
         mdir = getArgString args InputDir
         kdir = getArgString args AnnotationKeyInputDir
         sdat = getDataType args
     case (sdat, cStr, mf1,mdir, p, ky, kdir) of
       -- parse a string of chords
       (ChordLab, Just c, Nothing, Nothing , False,Nothing,Nothing)  ->
         do pr <- parseTreeVerb g o c 
            mapM_ (print . gTreeHead) (parsedPiece pr)
       -- and print a parsetree     
       (ChordLab, Just c, Nothing, Nothing , True,Nothing,Nothing)   ->
         do pr <- parseTree g o c 
            let ts = map gTreeHead (parsedPiece pr)
            _ <- printTreeHAn (pieceTreeHAn pr) (trimFilename ("pp" ++ c))
            printTreeHAnF ts (trimFilename c) >> return ()
       -- Parse one file, show full output
       (ChordLab, Nothing, Just f1, Nothing , False,Nothing,Nothing) -> 
         do pr  <- readFile f1 >>= parseTreeVerb g o
            print (pieceTreeHAn pr)
            mapM_ (print . gTreeHead) (parsedPiece pr)
       (ChordLab, Nothing, Just f1, Nothing , True ,Nothing,Nothing) ->
       --with post processing
         do pr <- readFile f1 >>= parseTree g o
            let ts = map gTreeHead (parsedPiece pr)
            printTreeHAn (pieceTreeHAn pr) (f1 ++ ".postProc") >> return ()
            printTreeHAnF ts f1 >> return () 
       -- Parse all files in one dir, show condensed output 
       (ChordLab, Nothing, Nothing, Just dir, False, Nothing, Nothing) ->
         parseDir g o dir bOut
       -- ** audio ground-truth annotation part **   
       (Annotation, Nothing, Just f1,Nothing, False, Just kf, Nothing) -> 
       -- parse a ground-truth annotation and its key and give verbose output      
         do key <- readFile kf
            pr  <- readFile f1 >>= parseAnnotationVerb g o key 
            print (pieceTreeHAn pr)
            mapM_ (print . gTreeHead) (take 10 $ parsedPiece pr) 
       (Annotation, Nothing, Just f1, Nothing , True , Just kf,Nothing) ->
       -- parse a ground-truth annotation and its key and print the parse 
         do key <- readFile kf
            pr  <- readFile f1 >>= parseAnnotation g o key 
            let ts = map gTreeHead (parsedPiece pr)
            printTreeHAn (pieceTreeHAn pr) (f1 ++ ".postProc") >> return ()
            printTreeHAnF ts f1 >> return ()   
       -- Parse all files in one dir, show condensed output 
       (Annotation, Nothing, Nothing, Just dir,False,Nothing,Just kd) ->
         parseAnnotationDir g o kd dir
       -- Else throw error
       _ -> usageError args err1
                   
trimFilename :: String -> String
trimFilename = filter (\x -> not (elem x ":*")) . concat . words . take 20 
                   
mainMatch :: Args MyArgs -> [PPOption] -> Bool -> MatchMode -> GrammarEx -> IO ()
mainMatch args o p m (GrammarEx g) =
     do let cStr  = getArgString args SourceInputString
            mf1   = getArg args SourceInputFile
            mf2   = getArg args TargetInputFile
            mdir  = getArg args InputDir
            bIn   = getArgString args BinaryIn
            me    = getArg args MaxErrorRate
        case (cStr,mf1,mf2,mdir,p) of
          -- Parse source and target file, show full output
          (_,Just f1, Just f2, Nothing, prnt)   -> 
            do c1 <- readFile' f1
               c2 <- readFile' f2
               matchFiles o m prnt c1 c2 f1 f2
          (Just c, Just f1, Nothing, Nothing, True) ->
            matchFiles o m True c f1 (trimFilename c) (trimFilename f1)
          -- match all files in one dir, show condensed output
          (_,Nothing, Nothing, Just dir, False) -> dirMatch g o bIn m me dir 
          _                                     -> usageError args err2

matchFiles :: [PPOption] -> MatchMode -> Bool -> String -> String 
           -> FilePath -> FilePath -> IO ()
matchFiles o m prnt f1 f2 n1 n2 = 
  -- should move to HarmTrace.IO.Main
  let (ParseResult key1 toks1 _ ts1 _nr1 te1 pe1 _) 
         = postProc o $ string2Piece Jazz f1
      (ParseResult key2 toks2 _ ts2 _nr2 te2 pe2 _) 
         = postProc o $ string2Piece Jazz f2
  in
  do  if not $ null te1 then showErrors "tokenizer 1: " te1 else putStr ""
      if not $ null te2 then showErrors "tokenizer 2: " te2 else putStr ""
      if not $ null pe1 then showErrors "parser 1: " pe1 else putStr ""
      if not $ null pe2 then showErrors "parser 2: " pe2 else putStr ""
      case (m,prnt) of
        (STDiff,_)      -> print (diffChordsLen toks1 toks2)
        (Align  ,False) -> print (getAlignDist key1 key2 toks1 toks2)
        (Align  ,True ) -> do let (mat,v,t) = alignChordLab key1 key2 toks1 toks2 
                              pPrintV t; print mat ; print v
                              
        (HAnAlign,True ) -> do let (mat,v,t) = alignHAnChord ts1 ts2
                               pPrintV t; print mat ; print v
        -- quick and dirty LCES plotting (should move to HarmTrace.IO)
        (LCESsize,True)  -> 
           do printTreeHAn ts1 (n1 ++ ".postProc") >> return ()
              printTreeHAn ts2 (n2 ++ ".postProc") >> return ()
              printTreeHAnF (fst $ getLCES ts1 ts2) 
                     (  (take 10 $ takeFileName n1) ++ ".vs." 
                     ++ (take 10 $ takeFileName n2) ++ ".lces") >> return ()
        _                -> error "Unimplemented." 

-- given a audio feature description files and groud-truth annotation, evaluates
-- a model-based chord labelling
mainChordRec :: (Maybe [TimedData Key] -> AudioFeat -> ChordAnnotation) 
             -> Args MyArgs -> Bool -> IO ()
mainChordRec ann args p = 
     do let af     = getArgString args SourceInputFile
            gt     = getArg args GroundTruthInputFile
            key    = getArg args SourceKeyInputFile
            afdir  = getArg args InputDir
            gtdir  = getArg args GroundTruthInputDir
            keydir = getArg args AnnotationKeyInputDir
        case (gt,af,key,gtdir,afdir,keydir,p) of
          -- evaluates a single audio feature set
          (Just g, Just a, Nothing, Nothing, Nothing, Nothing, prnt) -> 
            evaluateLabeling ann prnt g a Nothing >>= print 
          -- evaluates a single audio feature set and key annotation
          (Just g, Just a, k, Nothing, Nothing, Nothing, prnt) -> 
            evaluateLabeling ann prnt g a k >>= print   
          -- evaluates a directory with audio features
          (Nothing, Nothing, Nothing, Just gd, Just ad, Nothing, _ ) -> 
            batchLabeling ann gd ad Nothing
          -- evaluates a directory with audio features and key annotations
          (Nothing, Nothing, Nothing, Just gd, Just ad, k, _ ) -> 
            batchLabeling ann gd ad k
          _                                        -> usageError args err4