diff --git a/LICENSE b/LICENSE new file mode 100644 --- /dev/null +++ b/LICENSE @@ -0,0 +1,29 @@ +BSD 3-Clause License + +Copyright (c) 2020, Jiasen Wu +All rights reserved. + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are met: + +1. Redistributions of source code must retain the above copyright notice, this + list of conditions and the following disclaimer. + +2. Redistributions in binary form must reproduce the above copyright notice, + this list of conditions and the following disclaimer in the documentation + and/or other materials provided with the distribution. + +3. Neither the name of the copyright holder nor the names of its + contributors may be used to endorse or promote products derived from + this software without specific prior written permission. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. diff --git a/fei-datasets.cabal b/fei-datasets.cabal new file mode 100644 --- /dev/null +++ b/fei-datasets.cabal @@ -0,0 +1,101 @@ +cabal-version: 2.4 +name: fei-datasets +version: 1.0.0 +synopsis: Some datasets +description: +homepage: http://github.com/pierric/fei-datasets +license: BSD-3-Clause +license-file: LICENSE +author: Jiasen Wu +maintainer: jiasenwu@hotmail.com +copyright: 2020 - Jiasen Wu +category: Machine Learning, AI +build-type: Simple + +Library + exposed-modules: MXNet.NN.DataIter.PascalVOC + MXNet.NN.DataIter.Coco + MXNet.NN.DataIter.Anchor + MXNet.NN.DataIter.Common + hs-source-dirs: src + ghc-options: -Wall + default-language: Haskell2010 + default-extensions: GADTs, + TypeFamilies, + OverloadedLabels, + OverloadedLists, + OverloadedStrings, + FlexibleContexts, + FlexibleInstances, + StandaloneDeriving, + DeriveGeneric, + TypeOperators, + DataKinds, + PartialTypeSignatures + build-depends: base >= 4.7 && < 5.0 + , storable-tuple + , lens >= 4.12 + , transformers-base >= 0.4.4 + , aeson >= 1.2 + , repa >= 3.4 + , aeson >= 1.0 && <1.5 + , attoparsec (>=0.13.2.2 && <0.14) + , conduit >= 1.2 && < 1.4 + , hexpat + , store + , random-fu + , random-source + , conduit-concurrent-map + , resourcet + , hip + , rio + , vector + , fei-base + , fei-nn + , fei-cocoapi +Executable render + hs-source-dirs: utils + main-is: render.hs + default-language: Haskell2010 + build-depends: base >= 4.7 && < 5.0, + fei-base, + fei-nn, + fei-datasets, + optparse-applicative, + attoparsec, + resourcet, + lens, + conduit, + repa, + hip, + vector, + rio, + random-source, + JuicyPixels, + FontyFruity, + Rasterific + ghc-options: -threaded +Executable bench + hs-source-dirs: utils + main-is: bench.hs + default-language: Haskell2010 + build-depends: base >= 4.7 && < 5.0, + fei-base, + fei-nn, + fei-datasets, + optparse-applicative, + criterion, + attoparsec, + resourcet, + lens, + conduit, + repa, + hip, + rio, + random-source, + stm-conduit, + JuicyPixels, + JuicyPixels-repa, + JuicyPixels-extra + ghc-options: -threaded + diff --git a/src/MXNet/NN/DataIter/Anchor.hs b/src/MXNet/NN/DataIter/Anchor.hs new file mode 100644 --- /dev/null +++ b/src/MXNet/NN/DataIter/Anchor.hs @@ -0,0 +1,316 @@ +{-# LANGUAGE TemplateHaskell #-} +module MXNet.NN.DataIter.Anchor where + +import Control.Exception (throw) +import Control.Lens (ix, makeLenses, (^?!)) +import Data.Array.Repa ((:.) (..), All (..), Array, D, + DIM1, DIM2, U, Z (..), + fromListUnboxed, (+^)) +import qualified Data.Array.Repa as Repa +import Data.Random (StdRandom (..), runRVar, shuffleN) +import qualified Data.Vector.Unboxed.Mutable as UVM +import RIO +import qualified RIO.HashMap as M +import qualified RIO.NonEmpty as NE +import qualified RIO.Set as Set +import qualified RIO.Vector.Boxed as V +import qualified RIO.Vector.Boxed.Unsafe as V +import qualified RIO.Vector.Unboxed as UV +import qualified RIO.Vector.Unboxed.Partial as UV (maxIndex) +import qualified RIO.Vector.Unboxed.Unsafe as UV (unsafeFreeze) + +import MXNet.Base +import MXNet.Base.Operators.Tensor (__set_value, _slice) +import MXNet.Base.ParserUtils (decimal, list, parseR, rational, + tuple) +import MXNet.NN.Layer (copy, reshape) +import MXNet.NN.Utils.Repa (vstack, (^#!)) +import qualified MXNet.NN.Utils.Repa as Repa + +data AnchorError = BadDimension + deriving Show +instance Exception AnchorError + +type Anchor r = Array r DIM1 Float +type GTBox r = Array r DIM1 Float + +data Configuration = Configuration + { _conf_anchor_scales :: [Int] + , _conf_anchor_ratios :: [Float] + , _conf_anchor_base_size :: Int + , _conf_allowed_border :: Int + , _conf_fg_num :: Int + , _conf_batch_num :: Int + , _conf_bg_overlap :: Float + , _conf_fg_overlap :: Float + } + deriving Show +makeLenses ''Configuration + +anchors :: (Int, Int) -> Int -> Int -> [Int] -> [Float] -> V.Vector (Anchor U) +anchors (height, width) stride base_size scales ratios = + V.fromList + [ Repa.computeS $ anch +^ offs + | offY <- grid height + , offX <- grid width + , anch <- base + , let offs = fromListUnboxed (Z :. 4) [offX, offY, offX, offY]] + where + base = baseAnchors base_size scales ratios + grid size = map fromIntegral [0, stride .. size * stride-1] + +baseAnchors :: Int -> [Int] -> [Float] -> [Anchor U] +baseAnchors base_size scales ratios = [makeBase s r | r <- ratios, s <- scales] + where + makeBase :: Int -> Float -> Anchor U + makeBase scale ratio = + let sizeF = fromIntegral base_size - 1 + (w, h, x, y) = whctr (0, 0, sizeF, sizeF) + ws = round $ sqrt (w * h / ratio) :: Int + hs = round $ (fromIntegral ws) * ratio :: Int + in mkanchor x y (fromIntegral $ ws * scale) (fromIntegral $ hs * scale) + +whctr :: (Float, Float, Float, Float) -> (Float, Float, Float, Float) +whctr (x0, y0, x1, y1) = (w, h, x, y) + where + w = x1 - x0 + 1 + h = y1 - y0 + 1 + x = x0 + 0.5 * (w - 1) + y = y0 + 0.5 * (h - 1) + +mkanchor :: Float -> Float -> Float -> Float -> Anchor U +mkanchor x y w h = fromListUnboxed (Z :. 4) [x - hW, y - hH, x + hW, y + hH] + where + hW = 0.5 * (w - 1) + hH = 0.5 * (h - 1) + +(%!) :: HasCallStack => V.Vector a -> Int -> a +(%!) = (V.unsafeIndex) + +overlapMatrix :: Set Int -> V.Vector (GTBox U) -> V.Vector (Anchor U) -> Array D DIM2 Float +overlapMatrix goodIndices gtBoxes anBoxes = Repa.fromFunction (Z :. width :. height) calcOvp + where + width = V.length gtBoxes + height = V.length anBoxes + + calcArea box = (box ^#! 2 - box ^#! 0 + 1) * (box ^#! 3 - box ^#! 1 + 1) + areaA = V.map calcArea anBoxes + areaG = V.map calcArea gtBoxes + + calcOvp (Z :. ig :. ia) = + let gt = gtBoxes %! ig + anchor = anBoxes %! ia + iw = min (gt ^#! 2) (anchor ^#! 2) - max (gt ^#! 0) (anchor ^#! 0) + ih = min (gt ^#! 3) (anchor ^#! 3) - max (gt ^#! 1) (anchor ^#! 1) + areaI = iw * ih + areaU = areaA %! ia + areaG %! ig - areaI + in if Set.member ia goodIndices && iw > 0 && ih > 0 then areaI / areaU else 0 + +type Labels = Repa.Array U DIM2 Float -- UV.Vector Int +type Targets = Repa.Array U DIM2 Float -- UV.Vector (Float, Float, Float, Float) +type Weights = Repa.Array U DIM2 Float -- UV.Vector (Float, Float, Float, Float) + +assign :: (MonadReader Configuration m, MonadIO m) => + V.Vector (GTBox U) -> Int -> Int -> V.Vector (Anchor U) -> m (Labels, Targets, Weights) +assign gtBoxes imWidth imHeight anBoxes + | numGT == 0 = do + goodIndices <- filterGoodIndices + liftIO $ do + indices <- runRVar (shuffleN (Set.size goodIndices) (Set.toList goodIndices)) StdRandom + labels <- UVM.replicate numLabels (-1) + forM_ indices $ flip (UVM.write labels) 0 + let targets = UV.replicate (numLabels * 4) 0 + weights = UV.replicate (numLabels * 4) 0 + labels <- UV.unsafeFreeze labels + let labelsRepa = Repa.fromUnboxed (Z:.numLabels:.1) labels + targetsRepa = Repa.fromUnboxed (Z:.numLabels:.4) targets + weightsRepa = Repa.fromUnboxed (Z:.numLabels:.4) weights + return (labelsRepa, targetsRepa, weightsRepa) + + | otherwise = do + _fg_overlap <- view conf_fg_overlap + _bg_overlap <- view conf_bg_overlap + _batch_num <- view conf_batch_num + _fg_num <- view conf_fg_num + + goodIndices <- filterGoodIndices + + -- traceShowM ("#Good Anchors:", V.length goodIndices) + + liftIO $ do + -- TODO filter valid anchor boxes + labels <- UVM.replicate numLabels (-1) + + overlaps <- return $ Repa.computeUnboxedS $ overlapMatrix goodIndices gtBoxes anBoxes + -- for each GT, the hightest overlapping anchor is FG. + forM_ ([0..numGT-1] :: [_]) $ \i -> do + let s = Repa.computeUnboxedS $ slice overlaps 0 i + m = s ^#! argMax s + UV.mapM_ (flip (UVM.write labels) 1) $ UV.findIndices (==m) (Repa.toUnboxed s) + + -- FG anchors that have overlapping with any GT >= thresh + -- BG anchors that have overlapping with all GT < thresh + UV.forM_ (UV.indexed $ Repa.toUnboxed $ Repa.foldS max 0 $ Repa.transpose overlaps) $ \(i, m) -> do + when (Set.member i goodIndices) $ do + when (m >= _fg_overlap) $ do + -- traceShowM ("FG enable ", m, i) + (UVM.write labels i 1) + when (m < _bg_overlap) $ do + -- s <- UVM.read labels i + -- when (s == 1) $ traceShowM ("FG disable ", m, i) + (UVM.write labels i 0) + + -- subsample FG anchors if there are too many + fgs <- UV.findIndices (==1) <$> UV.unsafeFreeze labels + let numFG = UV.length fgs + when (numFG > _fg_num) $ do + indices <- runRVar (shuffleN numFG $ UV.toList fgs) StdRandom + -- traceShowM ("Disable A", take (numFG - _fg_num) indices) + forM_ (take (numFG - _fg_num) indices) $ + flip (UVM.write labels) (-1) + + -- subsample BG anchors if there are too many + bgs <- UV.findIndices (==0) <$> UV.unsafeFreeze labels + let numBG = UV.length bgs + maxBG = _batch_num - min numFG _fg_num + when (numBG > maxBG) $ do + indices <- runRVar (shuffleN numBG $ UV.toList bgs) StdRandom + -- traceShowM ("Disable B", take (numBG - maxBG) indices) + forM_ (take (numBG - maxBG) indices) $ + flip (UVM.write labels) (-1) + + -- compute the regression from each FG anchor to its gt + fgs <- UV.findIndices (==1) <$> UV.unsafeFreeze labels + let gts = UV.map (argMax . Repa.computeS . slice overlaps 1) fgs + gtDiffs = UV.zipWith makeTarget fgs gts + targets <- UVM.replicate numLabels (0, 0, 0, 0) + UV.zipWithM_ (UVM.write targets) fgs gtDiffs + + -- indicates which anchors have a regression + weights <- UVM.replicate numLabels (0, 0, 0, 0) + UV.forM_ fgs $ flip (UVM.write weights) (1, 1, 1, 1) + + labels <- UV.unsafeFreeze labels + targets <- UV.unsafeFreeze targets + weights <- UV.unsafeFreeze weights + let labelsRepa = Repa.fromUnboxed (Z:.numLabels:.1) labels + targetsRepa = Repa.fromUnboxed (Z:.numLabels:.4) (flattenT targets) + weightsRepa = Repa.fromUnboxed (Z:.numLabels:.4) (flattenT weights) + return (labelsRepa, targetsRepa, weightsRepa) + where + numGT = length gtBoxes :: Int + numLabels = length anBoxes :: Int + + -- + -- TODO: replace slice and argMax with the one from Utils + -- + slice :: _ -> Int -> Int -> _ + slice mat 0 ind = Repa.slice mat $ Z :. ind :. All + slice mat 1 ind = Repa.slice mat $ Z :. All :. ind + slice _ _ _ = throw BadDimension + + argMax :: Array U DIM1 Float -> Int + argMax = UV.maxIndex . Repa.toUnboxed + + asTuple :: Array U DIM1 Float -> (Float, Float, Float, Float) + asTuple box = (box ^#! 0, box ^#! 1, box ^#! 2, box ^#! 3) + + filterGoodIndices :: MonadReader Configuration m => m (Set Int) + filterGoodIndices = do + _allowed_border <- fromIntegral <$> view conf_allowed_border + let goodAnchor (x0, y0, x1, y1) = + x0 >= -_allowed_border && + y0 >= -_allowed_border && + x1 < fromIntegral imWidth + _allowed_border && + y1 < fromIntegral imHeight + _allowed_border + return $ Set.fromList $ V.toList $ V.findIndices (goodAnchor . asTuple) anBoxes + + makeTarget :: Int -> Int -> (Float, Float, Float, Float) + makeTarget fgi gti = + let fgBox = anBoxes %! fgi + gtBox = gtBoxes %! gti + (w1, h1, cx1, cy1) = whctr $ asTuple fgBox + (w2, h2, cx2, cy2) = whctr $ asTuple gtBox + dx = (cx2 - cx1) / (w1 + 1e-14) + dy = (cy2 - cy1) / (h1 + 1e-14) + dw = log (w2 / w1) + dh = log (h2 / h1) + in (dx, dy, dw, dh) + + flattenT :: UV.Vector (Float, Float, Float, Float) -> UV.Vector Float + flattenT = UV.concatMap (\(a,b,c,d) -> UV.fromList [a,b,c,d]) + +-- +-- Symbol for Anchor Generator +-- +data AnchorGeneratorProp = AnchorGeneratorProp + { _ag_ratios :: [Float] + , _ag_scales :: [Int] + , _ag_anchors_alloc :: NDArray Float + } +makeLenses ''AnchorGeneratorProp + +instance CustomOperationProp AnchorGeneratorProp where + prop_list_arguments _ = ["feature"] + prop_list_outputs _ = ["anchors"] + prop_list_auxiliary_states _ = [] + prop_infer_shape prop [feature_shape] = + let STensor [_, _, h, w] = feature_shape + num_scales = length (prop ^. ag_scales) + num_ratios = length (prop ^. ag_ratios) + num_anchs = num_scales * num_ratios * h * w + anchors_shape = STensor [1, num_anchs, 4] + in ([feature_shape], [anchors_shape], []) + prop_declare_backward_dependency _ _ _ _ = [] + + data Operation AnchorGeneratorProp = AnchorGenerator AnchorGeneratorProp + prop_create_operator prop _ _ = return (AnchorGenerator prop) + +instance CustomOperation (Operation AnchorGeneratorProp) where + forward (AnchorGenerator prop) [ReqWrite] [feature] [anchors] _ _ = do + -- :param: feature, shape of (1, C, H, W) + -- :param: anchors, shape of (1, N, 4), where N is number of anchors + + let alloc = prop ^. ag_anchors_alloc + + -- get the height, width of the feature (B,C,H,W) + [_,_,h,w] <- NE.toList <$> ndshape (NDArray feature :: NDArray Float) + let beg = [0,0,0,0] + end = [1,1,h,w] + ret <- prim _slice (#data := alloc .& #begin:= beg .& #end:= end .& Nil) + ret <- reshape [1,-1,4] ret + void $ copy ret (NDArray anchors) + + backward _ [ReqWrite] _ _ [in_grad_0] _ _ = do + -- type annotation is necessary, because only a general form + -- can be inferred. + let set_zeros = __set_value (#src := 0 .& Nil) :: TensorApply NDArrayHandle + void $ set_zeros (Just ([in_grad_0])) + +buildAnchorGenerator :: HasCallStack => [(Text, Text)] -> IO AnchorGeneratorProp +buildAnchorGenerator params = do + let allocV = anchors alloc_size stride base_size scales ratios + -- convert from `Vector (Array [4])` -> Array [1, 1, N, 4] + allocR = expandD0 $ expandD0 $ vstack $ V.map expandD0 allocV + num_scales = length scales + num_ratios = length ratios + (height, width) = alloc_size + + allocA <- makeEmptyNDArray [1, 1, height, width, 4*num_scales*num_ratios] contextCPU + copyFromRepa allocA allocR + + return $ AnchorGeneratorProp + { _ag_scales = scales + , _ag_ratios = ratios + , _ag_anchors_alloc = allocA + } + where + paramsM = M.fromList params + stride = parseR decimal $ paramsM ^?! ix "stride" + scales = parseR (list decimal) $ paramsM ^?! ix "scales" + ratios = parseR (list rational) $ paramsM ^?! ix "ratios" + base_size = parseR decimal $ paramsM ^?! ix "base_size" + alloc_size = parseR (tuple decimal) $ paramsM ^?! ix "alloc_size" + +expandD0 :: (Repa.Shape sh, UV.Unbox e) => Array U sh e -> Array U (sh :. Int) e +expandD0 = Repa.expandDim 0 diff --git a/src/MXNet/NN/DataIter/Coco.hs b/src/MXNet/NN/DataIter/Coco.hs new file mode 100644 --- /dev/null +++ b/src/MXNet/NN/DataIter/Coco.hs @@ -0,0 +1,277 @@ +{-# LANGUAGE TemplateHaskell #-} +{-# LANGUAGE UndecidableInstances #-} +{-# LANGUAGE ViewPatterns #-} +module MXNet.NN.DataIter.Coco( + module MXNet.NN.DataIter.Common, + Configuration(..), CocoConfig, conf_width, Coco(..), + classes, coco, + cocoImageList, cocoImages, cocoImagesBBoxes, cocoImagesBBoxesMasks, + loadImage, loadBoundingBoxes, loadMasks, + augmentWithBBoxes, +) where + +import Control.Lens (ix, makeLenses, (^?!)) +import qualified Data.Aeson as Aeson +import Data.Array.Repa ((:.) (..), Z (..), + fromListUnboxed) +import qualified Data.Array.Repa as Repa +import Data.Conduit +import qualified Data.Conduit.Combinators as C (yieldMany) +import qualified Data.Conduit.List as C +import qualified Data.Random as RND (runRVar, shuffleN, + stdUniform) +import Data.Random.Source.StdGen (StdGen) +import qualified Data.Store as Store +import qualified Data.Vector.Storable as SV (unsafeCast) +import GHC.Float (double2Float) +import GHC.Generics (Generic) +import qualified Graphics.Image as HIP +import qualified Graphics.Image.Interface as HIP +import qualified Graphics.Image.Interface.Repa as HIP +import RIO +import qualified RIO.ByteString as SBS +import RIO.Directory +import RIO.FilePath +import qualified RIO.Map as M +import qualified RIO.NonEmpty as RNE +import qualified RIO.Vector.Boxed as V +import qualified RIO.Vector.Storable as SV + +import MXNet.Coco.Mask +import MXNet.Coco.Types +import MXNet.NN.DataIter.Common + +classes :: V.Vector String +classes = V.fromList [ + "person", "bicycle", "car", "motorcycle", "airplane", "bus", "train", + "truck", "boat", "traffic light", "fire hydrant", "stop sign", + "parking meter", "bench", "bird", "cat", "dog", "horse", "sheep", + "cow", "elephant", "bear", "zebra", "giraffe", "backpack", "umbrella", + "handbag", "tie", "suitcase", "frisbee", "skis", "snowboard", + "sports ball", "kite", "baseball bat", "baseball glove", "skateboard", + "surfboard", "tennis racket", "bottle", "wine glass", "cup", "fork", + "knife", "spoon", "bowl", "banana", "apple", "sandwich", "orange", + "broccoli", "carrot", "hot dog", "pizza", "donut", "cake", "chair", + "couch", "potted plant", "bed", "dining table", "toilet", "tv", + "laptop", "mouse", "remote", "keyboard", "cell phone", "microwave", + "oven", "toaster", "sink", "refrigerator", "book", "clock", "vase", + "scissors", "teddy bear", "hair drier", "toothbrush"] + +data Coco = Coco FilePath String Instance (M.Map Int Int) + deriving Generic +instance Store.Store Coco + +data FileNotFound = FileNotFound String String + deriving Show +instance Exception FileNotFound + +data instance Configuration "coco" = CocoConfig { + _conf_coco :: Coco, + _conf_width :: Int, + _conf_mean :: (Float, Float, Float), + _conf_std :: (Float, Float, Float) +} + +makeLenses 'CocoConfig +type CocoConfig = Configuration "coco" + +instance HasDatasetConfig CocoConfig where + type DatasetTag CocoConfig = "coco" + datasetConfig = id + +instance ImageDataset "coco" where + imagesMean = conf_mean + imagesStdDev = conf_std + +cached :: Store.Store a => String -> IO a -> IO a +cached name action = do + createDirectoryIfMissing True "cache" + hitCache <- doesFileExist path + if hitCache then + SBS.readFile path >>= Store.decodeIO + else do + obj <- action + SBS.writeFile path (Store.encode obj) + return obj + where + path = "cache/" ++ name + +coco :: String -> String -> IO Coco +coco base datasplit = cached (datasplit ++ ".store") $ do + let annotationFile = base "annotations" ("instances_" ++ datasplit ++ ".json") + inst <- raiseLeft (FileNotFound annotationFile) $ Aeson.eitherDecodeFileStrict' annotationFile + let cat_to_classid = M.fromList $ V.toList $ V.map + (\cat -> (cat ^. odc_id, get_cat_classid (cat ^. odc_name))) + (inst ^. categories) + return $ Coco base datasplit inst cat_to_classid + + where + get_cat_classid (flip V.elemIndex classes -> Just index) = index + get_cat_classid _ = error "index not found in classes" + + +cocoImageList :: (MonadReader env m, HasDatasetConfig env, DatasetTag env ~ "coco", MonadIO m) + => IORef StdGen -> ConduitT () Image m () +cocoImageList rand_gen = do + Coco _ _ inst _ <- view (datasetConfig . conf_coco) + let all_images = inst ^. images + all_images_shuffle <- liftIO $ + RND.runRVar (RND.shuffleN (length all_images) (V.toList all_images)) rand_gen + C.yieldMany all_images_shuffle -- .| C.iterM (liftIO . print) + + +cocoImages :: (MonadReader env m, HasDatasetConfig env, DatasetTag env ~ "coco", MonadIO m) + => IORef StdGen -> ConduitT () (String, ImageTensor, ImageInfo) m () +cocoImages rand_gen = cocoImageList rand_gen .| C.mapM build + where + build image = do + let filename = image ^. img_file_name + (img, info) <- loadImage image + return $!! (filename, img, info) + + +cocoImagesBBoxes :: (MonadReader env m, HasDatasetConfig env, DatasetTag env ~ "coco", MonadIO m) + => IORef StdGen -> ConduitT () (String, ImageTensor, ImageInfo, GTBoxes) m () +cocoImagesBBoxes rand_gen = cocoImageList rand_gen .| C.mapM build .| C.catMaybes + where + build image = do + let filename = image ^. img_file_name + (img, info) <- loadImage image + mboxes <- loadBoundingBoxes image + case mboxes of + Nothing -> return Nothing + Just boxes -> return $!! Just (filename, img, info, boxes) + + +cocoImagesBBoxesMasks :: (MonadReader env m, HasDatasetConfig env, DatasetTag env ~ "coco", MonadIO m) + => IORef StdGen -> ConduitT () (String, ImageTensor, ImageInfo, GTBoxes, Masks) m () +cocoImagesBBoxesMasks rand_gen = cocoImageList rand_gen .| C.mapM build .| C.catMaybes + where + build image = do + let filename = image ^. img_file_name + (img, info) <- loadImage image + mboxes <- loadBoundingBoxes image + mmasks <- loadMasks image + case liftA2 (,) mboxes mmasks of + Nothing -> return Nothing + Just (boxes, masks) -> return $!! Just (filename, img, info, boxes, masks) + + +augmentWithBBoxes :: MonadIO m + => IORef StdGen + -> (String, ImageTensor, ImageInfo, GTBoxes) + -> m (String, ImageTensor, ImageInfo, GTBoxes) +augmentWithBBoxes rand_gen inp@(ident, img, info, bboxes) = liftIO $ do + do_flip <- RND.runRVar RND.stdUniform rand_gen + return $ + if not do_flip + then inp + else let shp@(Z :. c :. h :. w) = Repa.extent img + flip_img (Z :. c :. y :. x) = Z :. c :. y :. (w - x - 1) + img_flipped = Repa.computeUnboxedS $ Repa.backpermute shp flip_img img + w' = fromIntegral w + flip_box box = Repa.fromUnboxed (Z :. 5) $ case Repa.toUnboxed box of + -- flip, and rebuild the top-left, bottom-right coords + [x0, y0, x1, y1, sc] -> [w'-x1-1, y0, w'-x0-1, y1, sc] + boxes_flipped = V.map flip_box bboxes + in (ident, img_flipped, info, boxes_flipped) + + +getImageScale :: Image -> Int -> (Float, Int, Int) +getImageScale img size + | oriW >= oriH = (sizeF / oriW, floor (oriH * sizeF / oriW), size) + | otherwise = (sizeF / oriH, size, floor (oriW * sizeF / oriH)) + where + oriW = fromIntegral $ img ^. img_width + oriH = fromIntegral $ img ^. img_height + sizeF = fromIntegral size + + +loadImage :: (MonadReader env m, HasDatasetConfig env, DatasetTag env ~ "coco", MonadIO m) + => Image -> m (ImageTensor, ImageInfo) +loadImage img = do + Coco base datasplit _ _ <- view (datasetConfig . conf_coco) + width <- view (datasetConfig . conf_width) + + let imgFilePath = base datasplit img ^. img_file_name + imgRGB <- liftIO $ raiseLeft (FileNotFound imgFilePath) $ HIP.readImage imgFilePath + + let (scale, imgH, imgW) = getImageScale img width + imgInfo = fromListUnboxed (Z :. 3) [fromIntegral imgH, fromIntegral imgW, scale] + + imgResized = HIP.resize HIP.Bilinear HIP.Edge (imgH, imgW) (imgRGB :: HIP.Image HIP.VS HIP.RGB Double) + imgPadded = HIP.canvasSize (HIP.Fill $ HIP.PixelRGB 0.5 0.5 0.5) (width, width) imgResized + imgRepa = Repa.fromUnboxed (Z:.width:.width:.3) $ + SV.convert $ + SV.unsafeCast $ + HIP.toVector imgPadded + + imgEval <- transform $ Repa.map double2Float imgRepa + return (Repa.computeUnboxedS imgEval, imgInfo) + + +loadBoundingBoxes :: (MonadReader env m, HasDatasetConfig env, DatasetTag env ~ "coco", MonadIO m) + => Image -> m (Maybe GTBoxes) +loadBoundingBoxes img = do + Coco _ _ inst cat_to_classid <- view (datasetConfig . conf_coco) + size <- view (datasetConfig . conf_width) + + let imgAnns = V.filter (\ann -> ann ^. ann_image_id == imageId) (inst ^. annotations) + (scale, _, _) = getImageScale img size + gt_boxes = V.fromList $ catMaybes $ map (makeGTBox cat_to_classid scale) $ V.toList imgAnns + + return $ if V.null gt_boxes then Nothing else Just gt_boxes + where + imageId = img ^. img_id + width = img ^. img_width + height = img ^. img_height + makeGTBox cat_to_classid scale ann = + let (x0, y0, x1, y1) = cleanBBox (ann ^. ann_bbox) + classId = cat_to_classid ^?! ix (ann ^. ann_category_id) + in + if ann ^. ann_area > 0 && x1 > x0 && y1 > y0 + then Just $ fromListUnboxed (Z :. 5) [x0*scale, y0*scale, x1*scale, y1*scale, fromIntegral classId] + else Nothing + cleanBBox (x, y, w, h) = + let x0 = max 0 x + y0 = max 0 y + x1 = min (fromIntegral width - 1) (x0 + max 0 (w-1)) + y1 = min (fromIntegral height - 1) (y0 + max 0 (h-1)) + in (x0, y0, x1, y1) + + +loadMasks :: (MonadReader env m, HasDatasetConfig env, DatasetTag env ~ "coco", MonadIO m) + => Image -> m (Maybe Masks) +loadMasks img = do + Coco _ _ inst _ <- view (datasetConfig . conf_coco) + size <- view (datasetConfig . conf_width) + + let imgAnns = V.filter (\ann -> ann ^. ann_image_id == imageId) (inst ^. annotations) + (_, imgH, imgW) = getImageScale img size + masks <- V.mapM (getMask imgH imgW size) imgAnns + return $ if V.null masks then Nothing else Just masks + where + imageId = img ^. img_id + width = img ^. img_width + height = img ^. img_height + getMask upH upW size anno = liftIO $ do + crle <- case anno ^. ann_segmentation of + SegRLE cnts _ -> frUncompressedRLE cnts height width + SegPolygon (RNE.nonEmpty -> Just polys) -> + frPoly (RNE.map SV.fromList polys) height width + _ -> throwString "Cannot build CRLE" + crle <- merge crle False + mask <- decode crle + let -- always single channel, since we have merged the masks + -- also note that HIP uses image HxW + image = HIP.fromRepaArrayS $ + Repa.transpose $ + Repa.map (HIP.PixelY . (*255)) $ + Repa.reshape (Z :. height :. width) mask + imgResized = HIP.resize HIP.Bilinear HIP.Edge + (upH, upW) + image + imgPadded = HIP.canvasSize (HIP.Fill $ HIP.PixelY 0) + (size, size) + imgResized + return $ Repa.computeUnboxedS $ Repa.map (\(HIP.PixelY e) -> e) $ HIP.toRepaArray imgPadded diff --git a/src/MXNet/NN/DataIter/Common.hs b/src/MXNet/NN/DataIter/Common.hs new file mode 100644 --- /dev/null +++ b/src/MXNet/NN/DataIter/Common.hs @@ -0,0 +1,67 @@ +{-# LANGUAGE DataKinds #-} +module MXNet.NN.DataIter.Common where + +import Data.Array.Repa ((:.) (..), Array, D, DIM1, DIM2, DIM3, U, + Z (..), fromListUnboxed, (*^), (+^), (-^), + (/^)) +import qualified Data.Array.Repa as Repa +import GHC.TypeLits (Symbol) +import RIO +import qualified RIO.Vector.Boxed as V + +type ImageTensor = Array U DIM3 Float +type ImageInfo = Array U DIM1 Float +type GTBoxes = V.Vector (Array U DIM1 Float) +type Masks = V.Vector (Array U DIM2 Word8) + +data family Configuration (dataset :: Symbol) + +class ImageDataset (a :: Symbol) where + imagesMean :: Getting (Float, Float, Float) (Configuration a) (Float, Float, Float) + imagesStdDev :: Getting (Float, Float, Float) (Configuration a) (Float, Float, Float) + +class HasDatasetConfig env where + type DatasetTag env :: Symbol + datasetConfig :: Lens' env (Configuration (DatasetTag env)) + +-- transform HWC -> CHW +transform :: (HasDatasetConfig env, + ImageDataset (DatasetTag env), + MonadReader env m, + Repa.Source r Float) + => Array r DIM3 Float -> m (Array D DIM3 Float) +transform img = do + mean <- view (datasetConfig . imagesMean) + std <- view (datasetConfig . imagesStdDev) + let broadcast = Repa.extend (Repa.Any :. height :. width) + mean' = broadcast $ fromTuple mean + std' = broadcast $ fromTuple std + chnFirst = Repa.backpermute newShape (\ (Z :. c :. h :. w) -> Z :. h :. w :. c) img + return $ (chnFirst -^ mean') /^ std' + where + Z :. height :. width :. chn = Repa.extent img + newShape = Z:. chn :. height :. width + +-- transform CHW -> HWC +transformInv :: (ImageDataset s, MonadReader (Configuration s) m, Repa.Source r Float) => + Array r DIM3 Float -> m (Array D DIM3 Float) +transformInv img = do + mean <- view imagesMean + std <- view imagesStdDev + let broadcast = Repa.extend (Repa.Any :. height :. width) + mean' = broadcast $ fromTuple mean + std' = broadcast $ fromTuple std + addMean = img *^ std' +^ mean' + return $ Repa.backpermute newShape (\ (Z :. h :. w :. c) -> Z :. c :. h :. w) addMean + where + (Z :. chn :. height :. width) = Repa.extent img + newShape = Z :. height :. width :. chn + +fromTuple (a, b, c) = fromListUnboxed (Z :. (3 :: Int)) [a,b,c] + +raiseLeft :: (MonadThrow m, Exception e) => (a -> e) -> m (Either a b) -> m b +raiseLeft exc act = act >>= either (throwM . exc) return + +instance (Repa.Shape sh, Unbox e) => NFData (Array U sh e) where + rnf arr = Repa.deepSeqArray arr () + diff --git a/src/MXNet/NN/DataIter/PascalVOC.hs b/src/MXNet/NN/DataIter/PascalVOC.hs new file mode 100644 --- /dev/null +++ b/src/MXNet/NN/DataIter/PascalVOC.hs @@ -0,0 +1,131 @@ +{-# LANGUAGE TemplateHaskell #-} +{-# LANGUAGE UndecidableInstances #-} +module MXNet.NN.DataIter.PascalVOC ( + module MXNet.NN.DataIter.Common, + Configuration(..), VOCConfig, conf_width, + classes, vocMainImages, loadImageAndBBoxes +) where + +import Control.Exception (throw) +import Control.Lens (makeLenses) +import Data.Array.Repa ((:.) (..), Z (..), fromListUnboxed) +import qualified Data.Array.Repa as Repa +import Data.Conduit +import qualified Data.Conduit.Combinators as C (yieldMany) +import qualified Data.Conduit.List as C +import qualified Data.Random as RND (runRVar, shuffleN, + stdUniform) +import Data.Random.Source.StdGen (StdGen) +import qualified Data.Vector.Storable as SV (unsafeCast) +import GHC.Float (double2Float) +import qualified Graphics.Image as HIP +import qualified Graphics.Image.Interface as HIP +import RIO +import qualified RIO.ByteString as B +import RIO.FilePath +import qualified RIO.Text as T +import qualified RIO.Vector.Boxed as V +import qualified RIO.Vector.Storable as SV +import Text.XML.Expat.Proc +import Text.XML.Expat.Tree + +import MXNet.Base.ParserUtils (parseR, rational) +import MXNet.NN.DataIter.Common + +data Exc = FileNotFound String String + | CannotParseAnnotation String + deriving Show +instance Exception Exc + +classes :: V.Vector Text +classes = V.fromList [ + "__background__", -- always index 0 + "aeroplane", "bicycle", "bird", "boat", + "bottle", "bus", "car", "cat", "chair", + "cow", "diningtable", "dog", "horse", + "motorbike", "person", "pottedplant", + "sheep", "sofa", "train", "tvmonitor"] + +data instance Configuration "voc" = VOCConfig { + _conf_base_dir :: FilePath, + _conf_width :: Int, + _conf_mean :: (Float, Float, Float), + _conf_std :: (Float, Float, Float) +} +makeLenses 'VOCConfig + +type VOCConfig = Configuration "voc" + +instance HasDatasetConfig VOCConfig where + type DatasetTag VOCConfig = "voc" + datasetConfig = id + +instance ImageDataset "voc" where + imagesMean = conf_mean + imagesStdDev = conf_std + +vocMainImages :: (MonadReader env m, HasDatasetConfig env, DatasetTag env ~ "voc", MonadIO m) => + String -> IORef StdGen -> ConduitT () String m () +vocMainImages datasplit rand_gen = do + base <- view (datasetConfig . conf_base_dir) + let imageset = base "ImageSets" "Main" datasplit <.> "txt" + content <- readFileUtf8 imageset + let image_list = T.lines content + all_images <- liftIO $ RND.runRVar (RND.shuffleN (length image_list) image_list) rand_gen + C.yieldMany all_images .| C.map T.unpack + +loadImageAndBBoxes :: (MonadReader env m, HasDatasetConfig env, DatasetTag env ~ "voc", MonadIO m) => + String -> m (Maybe (String, ImageTensor, ImageInfo, GTBoxes)) +loadImageAndBBoxes ident = do + width <- view (datasetConfig . conf_width) + base <- view (datasetConfig . conf_base_dir) + + let imgFilePath = base "JPEGImages" ident <.> "jpg" + imgRGB <- liftIO $ raiseLeft (FileNotFound imgFilePath) $ + (HIP.readImageExact HIP.JPG imgFilePath) + + let (imgH, imgW) = HIP.dims (imgRGB :: HIP.Image HIP.VS HIP.RGB Double) + imgH_ = fromIntegral imgH + imgW_ = fromIntegral imgW + width_ = fromIntegral width + (scale, imgW', imgH') = if imgW >= imgH + then (width_ / imgW_, width, floor (imgH_ * width_ / imgW_)) + else (width_ / imgH_, floor (imgW_ * width_ / imgH_), width) + imgInfo = fromListUnboxed (Z :. 3) [fromIntegral imgH', fromIntegral imgW', scale] + + imgResized = HIP.resize HIP.Bilinear HIP.Edge (imgH', imgW') imgRGB + imgPadded = HIP.canvasSize (HIP.Fill $ HIP.PixelRGB 0.5 0.5 0.5) (width, width) imgResized + imgRepa = Repa.fromUnboxed (Z:.width:.width:.3) $ + SV.convert $ + SV.unsafeCast $ + HIP.toVector imgPadded + + let annoFilePath = base "Annotations" ident <.> "xml" + xml <- liftIO $ B.readFile annoFilePath + gtBoxes <- case parse' defaultParseOptions xml of + Left err -> throw (CannotParseAnnotation annoFilePath) err + Right root -> do + let objs = findElements "object" (root :: Node Text Text) + return $ V.fromList $ catMaybes $ map (makeGTBox scale) objs + + if V.null gtBoxes + then return Nothing + else do + imgEval <- transform $ Repa.map double2Float imgRepa + -- deepSeq the array so that the workload are well parallelized. + return $!! Just (ident, Repa.computeUnboxedS imgEval, imgInfo, gtBoxes) + where + makeGTBox scale node = do + className <- textContent <$> findElement "name" node + bndbox <- findElement "bndbox" node + xmin <- textContent <$> findElement "xmin" bndbox + xmax <- textContent <$> findElement "xmax" bndbox + ymin <- textContent <$> findElement "ymin" bndbox + ymax <- textContent <$> findElement "ymax" bndbox + classId <- V.elemIndex className classes + let x0 = parseR rational xmin + x1 = parseR rational xmax + y0 = parseR rational ymin + y1 = parseR rational ymax + return $ fromListUnboxed (Z :. 5) [x0*scale, y0*scale, x1*scale, y1*scale, fromIntegral classId] + diff --git a/utils/bench.hs b/utils/bench.hs new file mode 100644 --- /dev/null +++ b/utils/bench.hs @@ -0,0 +1,56 @@ +import qualified Codec.Picture as JUC +import qualified Codec.Picture.Extra as JUC +import qualified Codec.Picture.Repa as RPJ +import Control.Monad.Trans.Resource +import Criterion.Main +import qualified Data.Array.Repa as Repa +import Data.Conduit +import Data.Conduit.Async +import qualified Data.Conduit.List as C +import Data.Random.Source.StdGen (mkStdGen) +import qualified Graphics.Image as HIP +import RIO +import RIO.Directory +import RIO.FilePath + +import qualified MXNet.NN.DataIter.Coco as Coco + +main = do + home <- getHomeDirectory + let imgFilePath = home ".mxnet/datasets/coco/val2017/000000121242.jpg" + Right imgjuc <- liftIO (JUC.readImage imgFilePath) + Right imghip <- liftIO (HIP.readImage imgFilePath :: IO (Either String (HIP.Image HIP.VS HIP.RGB Double))) + + rand_gen <- newIORef $ mkStdGen 99 + cocoInst <- Coco.coco (home ".mxnet/datasets/coco") "val2017" + let cocoConf = Coco.CocoConfig cocoInst 1024 (0.5, 0.5, 0.5) (1, 1, 1) + iter1 = Coco.cocoImages rand_gen + iter2 = Coco.cocoImagesBBoxes rand_gen + + defaultMain + [ + bench "scale-img-juicy" $ nfIO $ + let img1 = JUC.convertRGB8 imgjuc + img2 = JUC.scaleBilinear 1024 1024 img1 + img3 = RPJ.imgData (RPJ.convertImage img2 :: RPJ.Img RPJ.RGB) + in Repa.computeUnboxedP img3 + + , bench "scale-img-hip" $ nfIO $ + let img2 = HIP.resize HIP.Bilinear HIP.Edge (1024, 1024) imghip + in return img2 + + , bench "img-iter" $ nfIO $ + runResourceT $ + flip runReaderT cocoConf $ + runConduit $ iter1 .| C.take 10 + + , bench "img-iter-async" $ nfIO $ + runResourceT $ + flip runReaderT cocoConf $ + runCConduit $ iter1 =$=& C.take 10 + + , bench "img-iter-and-load-gt" $ nfIO $ + runResourceT $ + flip runReaderT cocoConf $ + runConduit $ iter2 .| C.take 10 + ] diff --git a/utils/render.hs b/utils/render.hs new file mode 100644 --- /dev/null +++ b/utils/render.hs @@ -0,0 +1,138 @@ +{-# 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 +