Initial release0.0.0

7 files changed, 433 insertions, 0 deletions

diff --git a/src/IRE/Application.hs b/src/IRE/Application.hs
new file mode 100644
index 0000000..530e4a4
--- /dev/null
+++ b/src/IRE/Application.hs
@@ -0,0 +1,35 @@
+{-# LANGUAGE OverloadedStrings #-}
+
+module IRE.Application (
+  app
+) where
+
+import Control.Monad.Trans (liftIO)
+import Network.Wai (Application)
+import Network.Wai.Middleware.Static (addBase, hasPrefix, staticPolicy, (>->))
+import Network.Wai.Parse (FileInfo(..))
+import System.FilePath.Posix ((</>))
+import Web.Scotty ( ActionM, ScottyM, file, files, get,
+  json, middleware, post, scottyApp )
+
+import IRE.Application.YOLO (findItems)
+import IRE.YOLO (Detector)
+
+app :: FilePath -> Detector -> IO Application
+app rootDir net = scottyApp $ ire rootDir net
+
+ire :: FilePath -> Detector -> ScottyM ()
+ire rootDir net = do
+  middleware $ staticPolicy (hasPrefix "static" >-> addBase rootDir)
+  get "/" $ file (rootDir </> "webui.html")
+
+  post "/findItems" $ apiFindItems net
+
+
+apiFindItems :: Detector -> ActionM ()
+apiFindItems net = do
+  fs <- files
+  let fc = head [ fileContent fi | (_, fi) <- fs ]
+  resp <- liftIO $ findItems net fc
+  json resp
+
diff --git a/src/IRE/Application/YOLO.hs b/src/IRE/Application/YOLO.hs
new file mode 100644
index 0000000..1d6ab37
--- /dev/null
+++ b/src/IRE/Application/YOLO.hs
@@ -0,0 +1,16 @@
+module IRE.Application.YOLO (
+  findItems
+) where
+
+import qualified Data.ByteString.Lazy as LBS
+
+import IRE.YOLO (Detector, Item, detect)
+import qualified IRE.Logging as Log
+
+
+findItems :: Detector -> LBS.ByteString -> IO [Item]
+findItems d fc = do
+  r <- detect d 0.3 0.3 (LBS.toStrict fc)
+  Log.debug $ show r
+  return r
+
diff --git a/src/IRE/Config.hs b/src/IRE/Config.hs
new file mode 100644
index 0000000..c2b90d6
--- /dev/null
+++ b/src/IRE/Config.hs
@@ -0,0 +1,51 @@
+{-# LANGUAGE OverloadedStrings #-}
+
+module IRE.Config (
+  ConfigFile(..)
+, YOLO(..)
+, defaultConfig
+) where
+
+import Control.Applicative (empty)
+import Data.Aeson (FromJSON, parseJSON)
+import Data.Yaml (Value(Object), (.:), (.:?), (.!=))
+
+
+data ConfigFile = ConfigFile {
+  cfPort       :: Int
+, cfSocket     :: Maybe FilePath
+, cfYOLO       :: YOLO
+} deriving (Show)
+
+data YOLO = YOLO {
+  yoloCfg     :: FilePath
+, yoloWeights :: FilePath
+, yoloNames   :: FilePath
+} deriving (Show)
+
+
+defaultConfig :: ConfigFile
+defaultConfig = ConfigFile {
+    cfPort = 8080
+  , cfSocket = Nothing
+  , cfYOLO = YOLO {
+      yoloCfg = "yolo.cfg"
+    , yoloWeights = "yolo.weights"
+    , yoloNames = "yolo.names"
+    }
+  }
+
+instance FromJSON ConfigFile where
+  parseJSON (Object m) = ConfigFile <$>
+        m .:? "port"    .!= cfPort defaultConfig
+    <*> m .:? "socket"  .!= cfSocket defaultConfig
+    <*> m .:? "yolo"    .!= cfYOLO defaultConfig
+  parseJSON _ = empty
+
+instance FromJSON YOLO where
+  parseJSON (Object m) = YOLO <$>
+        m .: "cfg"
+    <*> m .: "weights"
+    <*> m .: "names"
+  parseJSON _ = empty
+
diff --git a/src/IRE/Logging.hs b/src/IRE/Logging.hs
new file mode 100644
index 0000000..3343f0f
--- /dev/null
+++ b/src/IRE/Logging.hs
@@ -0,0 +1,99 @@
+module IRE.Logging (
+  LogLevel(..)
+, debug
+, error
+, info
+, level
+, start
+, warn
+) where
+
+import Prelude hiding (error)
+
+import Control.Applicative (empty)
+import Control.Concurrent (forkIO)
+import Control.Concurrent.Chan (Chan, newChan, readChan, writeChan)
+import Control.Monad (forever, when)
+import Data.Aeson (FromJSON, ToJSON)
+import Data.Char (toLower)
+import Data.IORef (IORef, newIORef, readIORef, writeIORef)
+import System.IO (hPrint, stderr)
+import System.IO.Unsafe (unsafePerformIO)
+import Text.Read (readMaybe)
+import qualified Data.Aeson as JSON
+import qualified Data.Text as T
+
+start :: LogLevel -> IO ()
+start None = return ()
+start lvl = do
+  writeIORef logLevel lvl
+  ch <- readIORef chanRef
+  _ <- forkIO . forever $ readChan ch >>= hPrint stderr
+  return ()
+
+info :: String -> IO ()
+info = send . Message Info
+
+warn:: String -> IO ()
+warn = send . Message Warning
+
+error:: String -> IO ()
+error = send . Message Error
+
+debug :: String -> IO ()
+debug = send . Message Debug
+
+
+send :: Message -> IO ()
+send msg@(Message l _) = do
+  lvl <- level
+  when (l <= lvl) $ do
+    ch <- readIORef chanRef
+    writeChan ch msg
+
+{-# NOINLINE chanRef #-}
+chanRef :: IORef (Chan Message)
+chanRef = unsafePerformIO (newChan >>= newIORef)
+
+{-# NOINLINE logLevel #-}
+logLevel :: IORef LogLevel
+logLevel = unsafePerformIO (newIORef None)
+
+level :: IO LogLevel
+level = readIORef logLevel
+
+
+data LogLevel = None | Error | Warning | Info | Debug
+  deriving (Enum, Ord, Eq)
+
+instance Show LogLevel where
+  show None    = "NONE"
+  show Error   = "ERROR"
+  show Warning = "WARN"
+  show Info    = "INFO"
+  show Debug   = "DEBUG"
+
+instance Read LogLevel where
+  readsPrec _ s
+        | l == "none"  = [ (None,    "") ]
+        | l == "error" = [ (Error,   "") ]
+        | l == "warn"  = [ (Warning, "") ]
+        | l == "info"  = [ (Info,    "") ]
+        | l == "debug" = [ (Debug,   "") ]
+        | otherwise = [ ]
+        where l = map toLower s
+
+instance ToJSON LogLevel where
+  toJSON = JSON.String . T.pack . show
+
+instance FromJSON LogLevel where
+  parseJSON (JSON.String s) =
+    maybe (fail $ "unknown log level: " ++ show s) return (readMaybe . T.unpack $ s)
+  parseJSON _ = empty
+
+
+data Message = Message LogLevel String
+
+instance Show Message where
+  show (Message lvl str) = show lvl ++ ": " ++ str
+
diff --git a/src/IRE/Server.hs b/src/IRE/Server.hs
new file mode 100644
index 0000000..8dabc73
--- /dev/null
+++ b/src/IRE/Server.hs
@@ -0,0 +1,88 @@
+module IRE.Server (
+  server
+) where
+
+import Control.Exception.Base (throwIO, catch, bracket)
+import Data.Yaml.Include (decodeFileEither)
+import Data.Bits ((.|.))
+import Network.Socket (socket, setSocketOption, bind, listen, close,
+  maxListenQueue, getSocketName, inet_addr, Family(AF_UNIX, AF_INET),
+  SocketType(Stream), SocketOption(ReuseAddr), Socket, SockAddr(..))
+import Network.Wai.Handler.Warp (Port, defaultSettings, runSettingsSocket)
+import System.Exit (exitFailure)
+import System.IO (hPutStrLn, stderr)
+import System.IO.Error (isDoesNotExistError)
+import System.Posix.Files ( fileExist, groupReadMode, groupWriteMode,
+  ownerReadMode, ownerWriteMode, removeLink, setFileMode, socketMode )
+
+import IRE.Application (app)
+import IRE.Config (ConfigFile(..), defaultConfig)
+import IRE.YOLO (newDetector)
+import qualified IRE.Logging as Log
+
+type Listen = Either Port FilePath
+
+server :: FilePath -> FilePath -> IO ()
+server rootDir configFile = do
+  cf <- readConfigFile configFile
+  net <- newDetector (cfYOLO cf)
+  Log.start Log.Debug
+  let socketSpec =
+        case cfSocket cf of
+          Just p  -> Right p
+          Nothing -> Left $ cfPort cf
+  bracket
+    (createSocket socketSpec)
+    closeSocket
+    ( \sock -> do
+      listen sock maxListenQueue
+      runSettingsSocket defaultSettings sock =<< app rootDir net)
+
+
+readConfigFile :: FilePath -> IO ConfigFile
+readConfigFile f = do
+  e <- fileExist f
+  if e then do
+    r <- decodeFileEither f
+    case r of
+      Left ex -> do
+        hPutStrLn stderr $ "FATAL: " ++ f ++ ": " ++ show ex
+        exitFailure
+      Right cf -> return cf
+  else return defaultConfig
+
+
+createSocket :: Listen -> IO Socket
+createSocket (Right path) = do
+  removeIfExists path
+  sock <- socket AF_UNIX Stream 0
+  bind sock $ SockAddrUnix path
+  setFileMode path $ socketMode
+                  .|. ownerWriteMode .|. ownerReadMode
+                  .|. groupWriteMode .|. groupReadMode
+  Log.info $ "Listening on UNIX socket `" ++ path ++ "'"
+  return sock
+createSocket (Left port) = do
+  sock <- socket AF_INET Stream 0
+  setSocketOption sock ReuseAddr 1
+  addr <- inet_addr "127.0.0.1"
+  bind sock $ SockAddrInet (fromIntegral port) addr
+  Log.info $ "Listening on localhost:" ++ show port
+  return sock
+
+
+closeSocket :: Socket -> IO ()
+closeSocket sock = do
+  name <- getSocketName sock
+  close sock
+  case name of
+    SockAddrUnix path -> removeIfExists path
+    _ -> return ()
+
+
+removeIfExists :: FilePath -> IO ()
+removeIfExists fileName = removeLink fileName `catch` handleExists
+  where handleExists e
+          | isDoesNotExistError e = return ()
+          | otherwise = throwIO e
+
diff --git a/src/IRE/YOLO.hsc b/src/IRE/YOLO.hsc
new file mode 100644
index 0000000..9cec3c4
--- /dev/null
+++ b/src/IRE/YOLO.hsc
@@ -0,0 +1,97 @@
+{-# LANGUAGE OverloadedStrings #-}
+
+module IRE.YOLO (
+  Detector
+, Item(..)
+, detect
+, newDetector
+) where
+
+import Data.Aeson (ToJSON, object, toJSON, (.=))
+import Data.ByteString (ByteString)
+import Data.ByteString.Unsafe (unsafeUseAsCStringLen)
+import Data.Int (Int32)
+import Data.Text (Text)
+import Foreign.C.String (CString, withCString)
+import Foreign.C.Types (CChar, CFloat(..), CSize(..))
+import Foreign.Marshal.Array (peekArray)
+import Foreign.Ptr (Ptr)
+import Foreign.Storable (Storable(..))
+import qualified Control.Concurrent.Lock as L
+import qualified Data.Array.IArray as A
+import qualified Data.Text as T
+import qualified Data.Text.IO as TIO
+
+import IRE.Config (YOLO(..))
+
+#include "libdarknet.h"
+
+data Detector = Detector (Ptr ()) (A.Array Int Text) L.Lock
+
+data Item = Item {
+  itemClass      :: Int    -- ^ Object class number
+, itemName       :: Text   -- ^ human-readable description, e. g. "cat", "backpack".
+, itemConfidence :: Float  -- ^ A.K.A. probability
+, itemBox        :: (Float, Float, Float, Float)
+} deriving (Show)
+
+instance Storable Item where
+  sizeOf _ = #{size libdarknet_item}
+  alignment _ = #{alignment libdarknet_item}
+  poke = undefined
+  peek p = do
+    _c <- #{peek libdarknet_item, klass} p :: IO Int32
+    _p <- #{peek libdarknet_item, confidence} p
+    _x <- #{peek libdarknet_item, x} p
+    _y <- #{peek libdarknet_item, y} p
+    _h <- #{peek libdarknet_item, h} p
+    _w <- #{peek libdarknet_item, w} p
+    return $ Item (fromIntegral _c) "?" _p (_x, _y, _h, _w)
+
+instance ToJSON Item where
+  toJSON (Item c n a (x, y, h, w)) =
+    object
+      [ "class" .= c
+      , "name" .= n
+      , "confidence" .= a
+      , "box" .= object
+          [ "x" .= x
+          , "y" .= y
+          , "h" .= h
+          , "w" .= w
+          ]
+      ]
+
+newDetector :: YOLO -> IO Detector
+newDetector (YOLO cfg weights names) =
+  withCString cfg (\c ->
+    withCString weights (\w -> do
+      n <- T.lines <$> TIO.readFile names
+      let a = A.listArray (0, length n) n
+      l <- L.new
+      d <- libdarknet_new_detector c w
+      return $ Detector d a l))
+
+
+detect :: Detector -> Float -> Float -> ByteString -> IO [Item]
+detect (Detector d ns lk) threshold tree_threshold imgdata =
+  unsafeUseAsCStringLen imgdata (\(img, len) -> do
+    items <- L.with lk $ do
+      CSize s <- libdarknet_detect d
+                    (CFloat threshold) (CFloat tree_threshold)
+                    img (CSize $ fromIntegral len)
+      let ptr = libdarknet_get_items d
+      peekArray (fromIntegral s) ptr
+    return $ map (\i@(Item c _ _ _) -> i{itemName = ns A.! c}) items
+    )
+
+
+foreign import ccall safe "libdarknet_new_detector"
+  libdarknet_new_detector :: CString -> CString -> IO (Ptr ())
+
+foreign import ccall safe "libdarknet_detect"
+  libdarknet_detect :: Ptr () -> CFloat -> CFloat -> Ptr CChar -> CSize -> IO CSize
+
+foreign import ccall unsafe "libdarknet_get_items"
+  libdarknet_get_items :: Ptr () -> Ptr Item
+
diff --git a/src/Main.hs b/src/Main.hs
new file mode 100644
index 0000000..5070d8d
--- /dev/null
+++ b/src/Main.hs
@@ -0,0 +1,47 @@
+{-# LANGUAGE QuasiQuotes #-}
+module Main (
+  main
+) where
+
+import Data.Maybe (fromJust)
+import Data.Version (showVersion)
+import Paths_ire (getDataDir, version) -- from cabal
+import System.Environment (getArgs)
+import Text.InterpolatedString.Perl6 (qc)
+import qualified System.Console.Docopt.NoTH as O
+
+import IRE.Server (server)
+
+usage :: IO String
+usage = do
+  dataDir <- getDataDir
+  return $
+    "IRE " ++ showVersion version ++
+    " - Watching you!" ++ [qc|
+
+Usage:
+  ire [options]
+
+Options:
+  -c, --config=FILE        Configuration file [default: ire.yml]
+
+  -r, --rootdir=DIR        Web root directory with static files [default: {dataDir}]
+  -h, --help               Show this message
+
+Note:
+  The default configuration file is loaded if found,
+  otherwise default built-in settings are used.
+
+|]
+
+main :: IO ()
+main = do
+  doco <- O.parseUsageOrExit =<< usage
+  args <- O.parseArgsOrExit doco =<< getArgs
+  if args `O.isPresent` O.longOption "help"
+  then putStrLn $ O.usage doco
+  else do
+    let configFile = fromJust . O.getArg args $ O.longOption "config"
+        rootDir  = fromJust . O.getArg args $ O.longOption "rootdir"
+    server rootDir configFile
+