{-# LANGUAGE DataKinds #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RecordWildCards #-}
module Main where
import Codec.Picture.Types
import Control.Monad.Trans.Resource
import Data.Array.Repa ((:.) (..), Z (..))
import qualified Data.Array.Repa as Repa
import Data.Attoparsec.Text (char, decimal, endOfInput,
parseOnly, rational, sepBy)
import Data.Conduit
import qualified Data.Conduit.Combinators as C (map, mapM, mapM_, take)
import qualified Data.Conduit.List as C (catMaybes)
import Data.Random.Source.StdGen (mkStdGen)
import qualified Data.Vector.Storable as SV (unsafeCast)
import GHC.TypeLits (Symbol)
import qualified Graphics.Image as HIP
import qualified Graphics.Image.Interface as HIP
import Graphics.Rasterific
import Graphics.Rasterific.Texture (uniformTexture)
import Graphics.Text.TrueType
import Options.Applicative (Parser, auto, eitherReader,
execParser, fullDesc, header,
help, helper, info, long,
metavar, option, showDefault,
strOption, switch, value, (<**>))
import RIO
import RIO.FilePath
import qualified RIO.Text as T
import qualified RIO.Vector.Boxed as V
import qualified RIO.Vector.Boxed.Partial as V ((!))
import qualified RIO.Vector.Storable as SV
import qualified RIO.Vector.Unboxed as UV
import qualified MXNet.NN.DataIter.Coco as DC
import MXNet.NN.DataIter.Common
import qualified MXNet.NN.DataIter.PascalVOC as DV
data Args = Args
{ arg_dataset :: String
, arg_base_dir :: String
, arg_datasplit :: String
, arg_width :: Int
, arg_mean :: (Float, Float, Float)
, arg_stddev :: (Float, Float, Float)
, arg_num_imgs :: Int
, arg_shuffle :: Maybe Int
}
cmdArgParser = Args <$> strOption (long "dataset" <> metavar "DATASET" <> help "dataset name")
<*> strOption (long "base-dir" <> metavar "BASE" <> help "path to the root directory")
<*> strOption (long "datasplit" <> metavar "SPLIT" <> help "datasplit")
<*> option auto (long "img-size" <> metavar "SIZE" <> showDefault <> value 512 <> help "size of image")
<*> option floatList (long "img-pixel-means" <> metavar "RGB-MEAN" <> showDefault <> value (0.4850, 0.4580, 0.4076) <> help "RGB mean of images")
<*> option floatList (long "img-pixel-stddev" <> metavar "RGB-STD" <> showDefault <> value (1,1,1) <> help "RGB std-dev of images")
<*> option auto (long "num-imgs" <> metavar "NUM-IMG" <> showDefault <> value 10 <> help "number of images")
<*> option auto (long "shuffle" <> showDefault <> help "shuffle")
where
triple = do
a <- rational
char ','
b <- rational
char ','
c <- rational
endOfInput
return (a, b, c)
floatList = eitherReader (parseOnly triple . T.pack)
class HasWidth (a :: Symbol) where
targetWidth :: Getting Int (Configuration a) Int
instance HasWidth "voc" where
targetWidth = DV.conf_width
instance HasWidth "coco" where
targetWidth = DC.conf_width
renderWithBBox :: (HasWidth s, ImageDataset s, MonadReader (Configuration s) m, MonadIO m) =>
Font -> (String, V.Vector String, ImageTensor, ImageInfo, GTBoxes) -> m (String, HIP.Image HIP.VS HIP.RGBA HIP.Word8)
renderWithBBox font (ident, cls, img, info, gt) = do
width <- view targetWidth
let height = width
arr <- transformInv img
let rawUV = Repa.toUnboxed $ Repa.computeUnboxedS $ Repa.map (floor . (* 255.0)) arr :: UV.Vector HIP.Word8
rawSV = SV.unsafeCast $ UV.convert rawUV :: HIP.Vector HIP.VS (HIP.Pixel HIP.RGB HIP.Word8)
img = promoteImage $ HIP.toJPImageRGB8 $ HIP.fromVector (height, width) rawSV
res = renderDrawing width height (PixelRGBA8 0 0 0 0) $ do
drawImage img 0 (V2 0 0)
withTexture (uniformTexture $ PixelRGBA8 255 0 0 255) $ do
void $ forM (zip [0..] $ V.toList boxes) $ \(ind, [x0, y0, x1, y1, _]) -> do
stroke 1 JoinRound (CapRound, CapRound) $ rectangle (V2 x0 y0) (x1 - x0) (y1 - y0)
withTexture (uniformTexture $ PixelRGBA8 255 255 255 255) $ do
printTextAt font (PointSize 10) (V2 (x0+2) (y0+12)) (cls V.! ind)
return $ (ident, HIP.fromJPImageRGBA8 res)
where
boxes = V.map (UV.toList . Repa.toUnboxed) gt
lookupClassName :: V.Vector String -> (String, ImageTensor, ImageInfo, GTBoxes) -> (String, V.Vector String, ImageTensor, ImageInfo, GTBoxes)
lookupClassName table (imgname, tensor, info, gt) = (imgname, gtNames, tensor, info, gt)
where
gtNames = V.map ((table V.!) . floor . (`Repa.index` (Z:.4))) gt
main :: IO ()
main = do
fontCache <- buildCache
let fontPath = findFontInCache fontCache (FontDescriptor "Hack" (FontStyle False False))
fontPath <- case fontPath of
Nothing -> error "font not found"
Just a -> return a
font <- loadFontFile fontPath
font <- case font of
Left msg -> error msg
Right a -> return a
Args{..} <- execParser $ info (cmdArgParser <**> helper) fullDesc
let save (ident, img) = liftIO $ HIP.writeImageExact HIP.PNG [] (ident <.> "png") img
dump :: (HasWidth s, ImageDataset s, MonadReader (Configuration s) m, MonadIO m) =>
ConduitT (String, V.Vector String, ImageTensor, ImageInfo, GTBoxes) Void m ()
dump = C.take arg_num_imgs .|
C.mapM (renderWithBBox font) .|
C.mapM_ save
rand_gen <- newIORef $ mkStdGen $ fromMaybe 0 arg_shuffle
case arg_dataset of
"coco" -> do
coco <- DC.coco arg_base_dir arg_datasplit
let conf = DC.CocoConfig coco arg_width arg_mean arg_stddev
iter = DC.cocoImagesBBoxes rand_gen .| C.mapM (DC.augmentWithBBoxes rand_gen)
void $ runResourceT $ flip runReaderT conf $ runConduit $ iter .| C.map (lookupClassName DC.classes) .| dump
"voc" -> do
let conf = DV.VOCConfig arg_base_dir arg_width arg_mean arg_stddev
iter = DV.vocMainImages arg_datasplit rand_gen .| C.mapM DV.loadImageAndBBoxes .| C.catMaybes
void $ flip runReaderT conf $ runConduit $ iter .| C.map (lookupClassName $ V.map T.unpack DV.classes) .| dump