diff --git a/examples/ImageUtils.hs b/examples/ImageUtils.hs new file mode 100644 --- /dev/null +++ b/examples/ImageUtils.hs @@ -0,0 +1,194 @@ +{-# LANGUAGE ViewPatterns #-} +module Main where + +import qualified Codec.Picture as G +import Control.Lens (ix, (^?)) +import qualified Data.Aeson as Aeson +import Data.Array.Repa ((:.) (..), Z (..)) +import qualified Data.Array.Repa as RP +import Data.Colour.Palette.ColorSet (infiniteWebColors) +import Data.Colour.SRGB as COLOR +import qualified Data.Store as Store +import Formatting +import qualified Graphics.Image as G (Pixel (PixelY), VS (..), + displayImage, exchange, + fromJPImageRGBA8, + toJPImageY8) +import Graphics.Image.Interface.Repa (fromRepaArrayP) +import qualified Graphics.Rasterific as G +import qualified Graphics.Rasterific.Texture as G +import Graphics.Text.TrueType as G +import Options.Applicative +import RIO +import qualified RIO.ByteString as SBS +import qualified RIO.ByteString.Lazy as LBS +import RIO.Directory +import RIO.FilePath +import qualified RIO.HashMap as M +import qualified RIO.NonEmpty as RNE +import qualified RIO.Vector.Boxed as V +import qualified RIO.Vector.Storable as SV +import qualified RIO.Vector.Unboxed as UV + +import MXNet.Coco.Index +import MXNet.Coco.Mask +import MXNet.Coco.Types hiding (info) + + +data ArgSpec = ListImages + | DumpImage + { _arg_image_id :: Int + } + deriving Show + +argspec :: Parser (Maybe String, Maybe String, ArgSpec) +argspec = liftA3 (,,) + (option auto (metavar "BASEDIR" <> value Nothing <> help "path to coco")) + (option auto (metavar "SPLIT" <> value Nothing <> help "data split")) + (subparser (listImg <> dumpImg)) + where + listImg = command "list" $ info (pure ListImages) mempty + dumpImg = command "dump" $ info + (DumpImage <$> argument auto (metavar "IMAGE_ID" <> help "image id") <**> helper) mempty + + +main = do + (base_user, split_user, argspec) <- execParser $ info (argspec <**> helper) (fullDesc <> header "Coco Utility") + home <- getHomeDirectory + let base = fromMaybe (home ".mxnet/datasets/coco") base_user + split = fromMaybe "train2017" split_user + anno_filename = formatToString ("instances_" % string % ".json") split + runSimpleApp $ do + cocoinst <- liftIO $ readAnnotations $ base "annotations" anno_filename + case argspec of + ListImages -> listImages cocoinst + DumpImage{..} -> dumpImage cocoinst base split _arg_image_id + +listImages cocoinst = do + logInfo $ display ("ImageID Height Width Filename" :: Text) + forM_ (cocoinst ^. images) $ \image -> do + logInfo . display $ sformat + (right 15 ' ' % right 8 ' ' % right 8 ' ' % fitRight 80) + (image ^. img_id) (image ^. img_height) (image ^. img_width) (image ^. img_file_name) + +dumpImage cocoinst base split imgid = do + let annos = V.filter (\ann -> ann ^. ann_image_id == imgid) $ cocoinst ^. annotations + image = V.filter (\img -> img ^. img_id == imgid) $ cocoinst ^. images + cat_table = M.fromList $ V.toList $ + V.map (\cat -> (cat ^. odc_id, cat ^. odc_name)) $ cocoinst ^. categories + + case image V.!? 0 of + Nothing -> throwString "Invalid Image ID" + Just image -> do + img <- liftIO $ G.readImage (base split image ^. img_file_name) + case img of + Left msg -> throwString msg + Right img -> do + font_path <- liftIO $ G.findFontOfFamily "DejaVu Sans" (G.FontStyle False False) >>= + maybe (throwString "font not found") return + font <- liftIO $ G.loadFontFile font_path >>= either throwString return + + let width = image ^. img_width + height = image ^. img_height + white = G.PixelRGBA8 255 255 255 255 + headColor = G.PixelRGBA8 0 0x86 0xc1 128 + boxColor = G.PixelRGBA8 0 0x86 0xc1 255 + textColor = G.PixelRGBA8 255 0 0 255 + pointSize = G.PointSize 8 + + buildAnno anno color = liftIO $ do + let cat_id = anno ^. ann_category_id + cat_name = M.lookupDefault "???" cat_id cat_table + bbox = anno ^. ann_bbox + COLOR.RGB r g b = COLOR.toSRGB24 color + mask <- getMask anno width height + return $ [ + AnnoBoundingBox cat_name font pointSize textColor headColor bbox boxColor, + AnnoMask mask (G.PixelRGB8 r g b)] + + annotations <- fmap concat $ zipWithM buildAnno (V.toList annos) infiniteWebColors + + let rimg = G.renderDrawing width height white $ do + G.drawImage (G.convertRGBA8 img) 0 (G.V2 0 0) + mapM_ renderAnnotation annotations + outfile = formatToString (int % ".png") imgid + logInfo . display $ sformat ("Writing image: " % string) outfile + liftIO $ G.writePng outfile rimg + + -- forM_ annos $ \anno -> liftIO $ do + -- mask <- getMask anno width height + -- let outfile = formatToString (int % "_" % int % ".png") imgid (anno ^. ann_id) + -- G.writePng outfile mask + +readFromCache path = do + bs <- SBS.readFile path + Store.decodeIO bs + +readFromJson :: (MonadIO m, Aeson.FromJSON a) => FilePath -> m a +readFromJson path = do + bs <- LBS.readFile path + case Aeson.decode' bs of + Nothing -> error $ "cannot parse annotation file: " ++ path + Just inst -> return inst + +store path obj = do + SBS.writeFile path (Store.encode obj) + return obj + +readAnnotations :: (Store.Store a, Aeson.FromJSON a) + => FilePath -> IO a +readAnnotations path = + readFromCache cache_file `catchIO` + (\_ -> readFromJson path >>= store cache_file) + where + cache_file = "./instance.store" + +getMask :: Annotation -> Int -> Int -> IO (G.Image G.Pixel8) +getMask anno width height = 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 Z :. c :. w :. h = RP.extent mask + if (c > 1) + then (throwString "More than 1 channel") + else do + -- HIP uses image HxW + let image = RP.transpose $ RP.map (G.PixelY . (*255)) $ RP.reshape (Z :. w :. h) mask + return $ G.toJPImageY8 $ G.exchange G.VS $ fromRepaArrayP image + +data AnnoBuilder = AnnoBoundingBox + { _anno_text :: String + , _anno_text_font :: Font + , _anno_text_size :: PointSize + , _anno_text_fg :: G.PixelRGBA8 + , _anno_text_bg :: G.PixelRGBA8 + , _anno_box :: (Float, Float, Float, Float) + , _anno_box_fg :: G.PixelRGBA8 + } + | AnnoMask + { _anno_mask :: G.Image G.Pixel8 + , _anno_mask_fg :: G.PixelRGB8 + } + +renderAnnotation :: AnnoBuilder -> G.Drawing G.PixelRGBA8 () +renderAnnotation AnnoBoundingBox{..} = do + let (x1, y1, w, h) = _anno_box + BoundingBox tx1 ty1 tx2 ty2 _ = stringBoundingBox _anno_text_font 96 _anno_text_size _anno_text + title = G.rectangle (G.V2 (x1-1) (y1-ty2+ty1-1)) (tx2-tx1+2) (ty2-ty1+1) + textBase = G.V2 x1 y1 + rect = G.rectangle (G.V2 x1 y1) w h + G.withTexture (G.uniformTexture _anno_text_bg) $ do + G.fill title + G.withTexture (G.uniformTexture _anno_text_fg) $ do + G.printTextAt _anno_text_font _anno_text_size (G.V2 x1 (y1-2)) _anno_text + G.withTexture (G.uniformTexture _anno_box_fg) $ do + G.stroke 1 G.JoinRound (G.CapRound, G.CapRound) rect + +renderAnnotation AnnoMask{..} = do + let G.PixelRGB8 r g b = _anno_mask_fg + color_mask = G.pixelMap (\m -> G.PixelRGBA8 r g b (floor $ fromIntegral m * 0.6)) _anno_mask + G.drawImage color_mask 0 (G.V2 0 0) diff --git a/examples/Mask.hs b/examples/Mask.hs deleted file mode 100644 --- a/examples/Mask.hs +++ /dev/null @@ -1,87 +0,0 @@ -module Main where - -import qualified Data.ByteString.Lazy as BS -import qualified Data.ByteString as SBS -import Control.Lens ((^.), (^?), ix) -import qualified Data.Vector as V -import qualified Data.Vector.Storable as SV -import qualified Data.Vector.Unboxed as UV -import qualified Data.Aeson as Aeson -import qualified Data.Array.Repa as RP -import Data.Array.Repa ((:.)(..), Z(..)) -import qualified Data.Array.Repa.Repr.ForeignPtr as RF -import Codec.Picture as JP -import Codec.Picture.Repa -import qualified Data.Store as Store -import Control.Exception.Base - -import MXNet.Coco.Mask -import MXNet.Coco.Types -import MXNet.Coco.Index - -data Y8 - -class ToDynamicImage a where - toDynamicImage :: Img a -> DynamicImage - -instance ToDynamicImage Y8 where - toDynamicImage (Img arr0) = ImageY8 $ JP.Image w h (SV.unsafeFromForeignPtr0 (RF.toForeignPtr arr) (h*w*z) ) - where - (Z :. h :. w :. z) = RP.extent arr - arr = RP.computeS arr0 - -readFromCache path = do - bs <- SBS.readFile path - Store.decodeIO bs - -readFromJson path = do - bs <- BS.readFile path - case Aeson.decode' bs of - Nothing -> error $ "cannot parse annotation file: " ++ path - Just inst -> return inst - -store path obj = do - SBS.writeFile path (Store.encode obj) - return obj - -readAnnotations path = - readFromCache cache_file `catch` (\ e -> do - let _ = e :: IOException - readFromJson path >>= - store cache_file) - where - cache_file = "./instance.store" - -annotatino_file = "/home/jiasen/dschungel/coco/annotations/instances_train2017.json" - -main = do - inst <- readAnnotations annotatino_file - mapM_ (\cat -> putStrLn $ cat ^. odc_name) $ allCats inst - let anno = V.head $ allAnns inst - imgId = anno ^. ann_image_id - img = V.head $ V.filter (\img -> img ^. img_id == imgId) (inst ^. images) - height = img ^. img_height - width = img ^. img_width - - store "./imgs.store" $ inst ^. images - store "./anns.store" $ inst ^. annotations - store "./cats.store" $ inst ^. categories - - - -- putStrLn $ show imgId - - -- crle <- case anno ^. ann_segmentation of - -- SegRLE cnts _ -> frUncompressedRLE cnts height width - -- SegPolygon polys -> frPoly (map SV.fromList polys) height width - - -- mask <- decode crle - - -- let Z :. c :. w :. h = RP.extent mask - -- maskHW = RP.backpermute (Z :. h :. w :. c) (\ (Z :. c :. w :. h) -> Z :. h :. w :. c) mask - -- maskImg = toDynamicImage $ (Img $ RP.map (*255) maskHW :: Img Y8) - - -- savePngImage "a.png" maskImg - - -- putStrLn $ img ^. img_file_name - -- putStrLn $ img ^. img_flickr_url - -- putStrLn $ img ^. img_coco_url diff --git a/examples/Profiling.hs b/examples/Profiling.hs deleted file mode 100644 --- a/examples/Profiling.hs +++ /dev/null @@ -1,23 +0,0 @@ -import Criterion.Main -import Criterion.Main.Options -import Data.Store -import qualified Data.IntSet as Set -import qualified Data.ByteString as BS -import qualified Data.Vector as V -import qualified Data.Vector.Unboxed as UV -import Data.Array.Repa ((:.)(..), Z (..), fromUnboxed, computeUnboxedP, computeUnboxedS) - -import MXNet.NN.DataIter.Anchor - -main = do - goodIndices <- BS.readFile "examples/goodIndices.bin" >>= decodeIO :: IO (V.Vector Int) - gtBoxes <- BS.readFile "examples/gtBoxes.bin" >>= decodeIO :: IO (V.Vector (UV.Vector Float)) - anchors <- BS.readFile "examples/anchors.bin" >>= decodeIO :: IO (V.Vector (UV.Vector Float)) - goodIndices <- return $ Set.fromList $ V.toList goodIndices :: IO Set.IntSet - gtBoxes <- return $ V.map (fromUnboxed (Z:.(5::Int))) gtBoxes - anchors <- return $ V.map (fromUnboxed (Z:.(4::Int))) anchors - - defaultMain - [ bench "computeUnboxedP" $ whnfIO $ computeUnboxedP $ overlapMatrix goodIndices gtBoxes anchors - , bench "computeUnboxedS" $ whnf computeUnboxedS $ overlapMatrix goodIndices gtBoxes anchors - ] diff --git a/fei-cocoapi.cabal b/fei-cocoapi.cabal --- a/fei-cocoapi.cabal +++ b/fei-cocoapi.cabal @@ -1,24 +1,22 @@ +cabal-version: 2.4 name: fei-cocoapi -version: 0.2.0 +version: 1.0.0 synopsis: Cocodataset with cocoapi description: Haskell binding for the cocoapi in c homepage: http://github.com/pierric/fei-cocoapi -license: BSD3 +license: BSD-3-Clause license-file: LICENSE author: Jiasen Wu maintainer: jiasenwu@hotmail.com -copyright: Copyright: (c) 2019 Jiasen Wu +copyright: 2020 - Jiasen Wu category: Machine Learning, AI build-type: Simple -cabal-version: 1.24 extra-source-files: cbits/*.h, cbits/*.c Library exposed-modules: MXNet.Coco.Types MXNet.Coco.Mask MXNet.Coco.Index - MXNet.NN.DataIter.Coco - MXNet.NN.DataIter.Anchor other-modules: MXNet.Coco.Raw hs-source-dirs: src ghc-options: -Wall @@ -27,63 +25,107 @@ TypeFamilies, OverloadedLabels, FlexibleContexts, + FlexibleInstances, StandaloneDeriving, DeriveGeneric, - TypeOperators + TypeOperators, + OverloadedStrings, + LambdaCase, + MultiWayIf, + DoAndIfThenElse, + TypeApplications, + DataKinds, + ExplicitForAll, + NoImplicitPrelude build-depends: base >= 4.7 && < 5.0 , storable-tuple - , vector >= 0.12 - , mtl >= 2.2 - , lens >= 4.12 - , transformers-base >= 0.4.4 - , aeson >= 1.2 - , containers >= 0.5 - , bytestring >= 0.10 - , exceptions >= 0.8.3 + , lens + , transformers-base + , aeson , time < 2.0 - , repa >= 3.4 - , JuicyPixels - , JuicyPixels-repa - , JuicyPixels-extra - , aeson >= 1.0 && <1.5 - , attoparsec (>=0.13.2.2 && <0.14) - , lens >= 4.12 - , conduit >= 1.2 && < 1.4 + , repa + , aeson + , attoparsec + , lens + , conduit , store - , filepath - , directory , random-fu - , fei-base - , fei-dataiter - Build-tools: c2hs + , conduit-concurrent-map + , unliftio-core + , resourcet + , hip + , rio + , vector + build-tool-depends: c2hs:c2hs c-sources: cbits/maskApi.c include-dirs: cbits/ includes: maskApi.h -Executable mask +Executable imageutils hs-source-dirs: examples - main-is: Mask.hs + main-is: ImageUtils.hs default-language: Haskell2010 + default-extensions: RecordWildCards, + OverloadedStrings build-depends: base >= 4.7 && < 5.0, - fei-cocoapi, - bytestring, + rio, lens, aeson, - vector, + optparse-applicative, JuicyPixels, - JuicyPixels-repa, repa, - store - -Executable profiling - hs-source-dirs: examples - main-is: Profiling.hs - default-language: Haskell2010 - build-depends: base >= 4.7 && < 5.0, - fei-cocoapi, - criterion, store, - repa, - bytestring, - vector, - containers + hip, + Rasterific, + FontyFruity, + formatting, + palette, + colour, + fei-cocoapi + +-- Executable profiling +-- hs-source-dirs: examples +-- main-is: profiling.hs +-- default-language: Haskell2010 +-- build-depends: base >= 4.7 && < 5.0, +-- fei-cocoapi, +-- fei-nn, +-- fei-base, +-- criterion, +-- store, +-- repa, +-- bytestring, +-- vector, +-- containers, +-- conduit, +-- resourcet, +-- lens, +-- deepseq, +-- mtl, +-- JuicyPixels, +-- JuicyPixels-repa, +-- JuicyPixels-extra, +-- hip, +-- conduit-concurrent-map +-- ghc-options: -threaded + +-- Executable hist +-- hs-source-dirs: examples +-- main-is: hist.hs +-- default-language: Haskell2010 +-- build-depends: base >= 4.7 && < 5.0, +-- fei-cocoapi, +-- fei-base, +-- fei-nn, +-- mtl, +-- store, +-- repa, +-- conduit, +-- lens, +-- bytestring, +-- vector, +-- containers, +-- resourcet, +-- histogram-fill, +-- conduit-concurrent-map +-- ghc-options: -threaded diff --git a/src/MXNet/Coco/Index.hs b/src/MXNet/Coco/Index.hs --- a/src/MXNet/Coco/Index.hs +++ b/src/MXNet/Coco/Index.hs @@ -1,28 +1,28 @@ module MXNet.Coco.Index where -import Control.Lens ((^.)) -import qualified Data.Vector as V (Vector, filter, null, head) +import RIO hiding (Category) +import qualified Data.Vector as V (filter, null, head) import MXNet.Coco.Types -allCats :: Instance -> V.Vector Category +allCats :: Instance -> Vector Category allCats = (^. categories) -allAnns :: Instance -> V.Vector Annotation +allAnns :: Instance -> Vector Annotation allAnns = (^. annotations) -catByName :: String -> V.Vector Category -> Maybe Category +catByName :: String -> Vector Category -> Maybe Category catByName name = vecToMaybe . V.filter (\cat -> cat ^. odc_name == name) -annsByCat :: Category -> V.Vector Annotation -> V.Vector Annotation +annsByCat :: Category -> Vector Annotation -> Vector Annotation annsByCat cat = V.filter (\ann -> ann ^. ann_category_id == cat ^. odc_id ) -annsByImg :: Image -> V.Vector Annotation -> V.Vector Annotation +annsByImg :: Image -> Vector Annotation -> Vector Annotation annsByImg img = V.filter (\ann -> ann ^. ann_image_id == img ^. img_id ) -annByCatImg :: Image -> Category -> V.Vector Annotation -> Maybe Annotation +annByCatImg :: Image -> Category -> Vector Annotation -> Maybe Annotation annByCatImg img cat = vecToMaybe . annsByCat cat . annsByImg img -vecToMaybe :: V.Vector a -> Maybe a +vecToMaybe :: Vector a -> Maybe a vecToMaybe vec | V.null vec = Nothing | otherwise = Just $ V.head vec diff --git a/src/MXNet/Coco/Mask.hs b/src/MXNet/Coco/Mask.hs --- a/src/MXNet/Coco/Mask.hs +++ b/src/MXNet/Coco/Mask.hs @@ -1,23 +1,27 @@ {-# LANGUAGE TypeOperators #-} module MXNet.Coco.Mask where -import Data.Word -import qualified Data.ByteString as BS -import Data.Array.Repa (Array, Z(..), (:.)(..), DIM1, DIM2, DIM3, extent) -import Data.Array.Repa.Repr.Unboxed -import qualified Data.Vector.Unboxed as UV (convert, length) -import qualified Data.Vector.Storable as SV (Vector, map, unsafeCast) +import Data.Array.Repa ((:.) (..), Array, DIM1, DIM2, + DIM3, Z (..), extent) +import Data.Array.Repa.Repr.Unboxed +import qualified Data.Vector.Storable as SV (unsafeCast) +import RIO +import qualified RIO.ByteString as BS +import qualified RIO.Vector.Storable as SV (Vector, map) +import qualified RIO.Vector.Unboxed as UV (convert, length) -import MXNet.Coco.Raw +import MXNet.Coco.Raw +-- NOTE: -- mask should be of 3 dimension and in CWH order +-- also assuming that CDouble has the same memory representation as Double type Mask = Array U DIM3 Word8 type Area = Array U DIM1 Word32 type Iou = Array U DIM2 Double type BBox = Array U DIM2 Double type Poly = SV.Vector Double -data CompactRLE = CompactRLE Int Int [BS.ByteString] +data CompactRLE = CompactRLE Int Int (NonEmpty BS.ByteString) encode :: Mask -> IO CompactRLE encode mask = do @@ -31,7 +35,7 @@ let n = length bss rles <- frString im raw <- rleDecode rles h w - return $ + return $ fromUnboxed (Z :. n :. w :. h) $ UV.convert $ SV.map fromIntegral raw @@ -44,8 +48,8 @@ rles <- frString im orle <- rleMerge rles intersect bs <- rleToString orle - return $ CompactRLE h w [bs] - else + return $ CompactRLE h w (bs :| []) + else return im area :: CompactRLE -> IO Area @@ -54,7 +58,7 @@ rles <- frString im as <- rleArea rles num -- assuming the area can be represented as Word32 - return $ + return $ fromListUnboxed (Z :. num) $ map fromIntegral $ as @@ -76,7 +80,7 @@ toBBox im = do rles <- frString im BB bb <- rleToBbox rles - let bb' = UV.convert $ SV.unsafeCast bb + let bb' = UV.convert $ SV.unsafeCast bb return $ fromUnboxed (Z :. UV.length bb' :. 4) bb' frBBox :: BBox -> Int -> Int -> IO CompactRLE @@ -84,7 +88,7 @@ rles <- rleFrBbox (BB $ SV.unsafeCast $ UV.convert $ toUnboxed bb) h w CompactRLE h w <$> mapM rleToString rles -frPoly :: [Poly] -> Int -> Int -> IO CompactRLE +frPoly :: NonEmpty Poly -> Int -> Int -> IO CompactRLE frPoly polys h w = do rles <- mapM (\poly -> rleFrPoly (SV.unsafeCast poly) h w) polys CompactRLE h w <$> mapM rleToString rles @@ -93,7 +97,7 @@ frUncompressedRLE raw h w = do orle <- rleInit h w (map fromIntegral raw) crle <- rleToString orle - return $ CompactRLE h w [crle] + return $ CompactRLE h w (crle :| []) -frString :: CompactRLE -> IO [RLE] +frString :: CompactRLE -> IO (NonEmpty RLE) frString (CompactRLE h w bss) = mapM (\bs -> rleFrString bs h w) bss diff --git a/src/MXNet/Coco/Raw.chs b/src/MXNet/Coco/Raw.chs --- a/src/MXNet/Coco/Raw.chs +++ b/src/MXNet/Coco/Raw.chs @@ -1,7 +1,14 @@ {-# LANGUAGE ScopedTypeVariables #-} module MXNet.Coco.Raw where -import Foreign.Storable +import RIO +import RIO.Partial (toEnum) +import qualified RIO.NonEmpty as RNE +import qualified RIO.NonEmpty.Partial as RNE +import qualified RIO.Vector.Storable as SV +import qualified RIO.Vector.Storable.Unsafe as SV +import qualified Data.Vector.Storable.Mutable as SVM +import qualified RIO.ByteString as BS import Foreign.Ptr import Foreign.ForeignPtr import Foreign.ForeignPtr.Unsafe (unsafeForeignPtrToPtr) @@ -10,10 +17,6 @@ import Foreign.Marshal.Array import Foreign.Marshal.Alloc import Foreign.Storable.Tuple () -import qualified Data.Vector.Storable as SV -import qualified Data.Vector.Storable.Mutable as SVM -import qualified Data.ByteString as BS -import Control.Exception #include "maskApi.h" @@ -25,12 +28,16 @@ } makeRLE :: (Ptr () -> IO ()) -> IO RLE -makeRLE a = makeRLEs 1 a >>= return . head +makeRLE a = makeRLEs 1 a >>= return . RNE.head -makeRLEs :: Int -> (Ptr () -> IO ()) -> IO [RLE] -makeRLEs num a = allocaBytesAligned (num * {#sizeof RLE #}) {#alignof RLE#} (\prle -> do - a prle - go num prle []) +makeRLEs :: Int -> (Ptr () -> IO ()) -> IO (NonEmpty RLE) +makeRLEs num a + | num < 1 = error "number should be positive" + | otherwise = do + rles <- allocaBytesAligned (num * {#sizeof RLE #}) {#alignof RLE#} (\prle -> do + a prle + go num prle []) + return $ RNE.fromList rles where go 0 _ rles = return $ reverse rles go n prle rles = do @@ -46,19 +53,19 @@ return $ RLE h w m mgr_c withRLE :: RLE -> (Ptr () -> IO a) -> IO a -withRLE rle = withRLEs [rle] +withRLE rle = withRLEs (RNE.fromList [rle]) -withRLEs :: [RLE] -> (Ptr () -> IO a) -> IO a +withRLEs :: NonEmpty RLE -> (Ptr () -> IO a) -> IO a withRLEs rles = withRLEsLen (length rles) rles -withRLEsLen :: Int -> [RLE] -> (Ptr () -> IO a) -> IO a +withRLEsLen :: Int -> NonEmpty RLE -> (Ptr () -> IO a) -> IO a withRLEsLen num rles a = do allocaBytesAligned (num * {#sizeof RLE#}) {#alignof RLE#} $ \prles -> do - go prles rles + go prles $ RNE.toList rles ret <- a prles mapM_ (touchForeignPtr . _rle_cnts) rles return ret - where + where go _ [] = return () go prles (rle : nrles) = do pokeRLE prles rle @@ -106,13 +113,15 @@ `Int' } -> `()' #} - -rleEncode :: SV.Vector CUChar -> Int -> Int -> Int -> IO [RLE] -rleEncode m h w n = do - makeRLEs n (\ prle -> - svUnsafeWith m (\pm -> do - rleEncode_ prle (castPtr pm) h w n)) +rleEncode :: SV.Vector CUChar -> Int -> Int -> Int -> IO (NonEmpty RLE) +rleEncode m h w n + | n < 1 = error "number must be positive" + | otherwise = do + makeRLEs n (\ prle -> + svUnsafeWith m (\pm -> do + rleEncode_ prle (castPtr pm) h w n)) + {#fun rleDecode as rleDecode_ { `Ptr ()', @@ -121,9 +130,9 @@ } -> `()' #} -rleDecode :: [RLE] -> Int -> Int -> IO (SV.Vector CUChar) +rleDecode :: NonEmpty RLE -> Int -> Int -> IO (SV.Vector CUChar) rleDecode rles h w = do - let n = length rles + let n = length rles size = n * h * w mv <- SVM.new size SVM.unsafeWith mv $ (\ptr -> do @@ -140,27 +149,29 @@ } -> `()' #} -rleMerge :: [RLE] -> Bool -> IO RLE +rleMerge :: NonEmpty RLE -> Bool -> IO RLE rleMerge rles intersect = do let num = length rles - withRLEsLen num rles $ \prles -> + withRLEsLen num rles $ \prles -> makeRLE $ \porle -> rleMerge_ prles porle num intersect {#fun rleArea as rleArea_ { - withRLEs* `[RLE]', + withRLEs* `NonEmpty RLE', `Int', id `Ptr CUInt' } -> `()' #} -rleArea :: [RLE] -> Int -> IO [CUInt] -rleArea r n = do +rleArea :: NonEmpty RLE -> Int -> IO [CUInt] +rleArea r n + | n < 1 = error "number must be positive" + | otherwise = do allocaArray n (\pa -> do rleArea_ r n pa peekArray n pa) - + {#fun rleIou as rleIou_ { `Ptr ()', @@ -172,28 +183,28 @@ } -> `()' #} -rleIou :: [RLE] -> [RLE] -> [Bool] -> IO ((Int,Int), [Double]) +rleIou :: NonEmpty RLE -> NonEmpty RLE -> [Bool] -> IO ((Int,Int), [Double]) rleIou dt gt iscrowd = do let m = length dt n = length gt c = length iscrowd - assert (n == c) $ allocaArray (m*n) $ \po -> - withRLEsLen m dt $ \pdt -> - withRLEsLen n gt $ \pgt -> do + assert (n == c) $ allocaArray (m*n) $ \po -> + withRLEsLen m dt $ \pdt -> + withRLEsLen n gt $ \pgt -> do rleIou_ pdt pgt m n (SV.fromList $ map (toEnum . fromEnum) iscrowd) po raw <- peekArray (m * n) po return $ ((m,n), map realToFrac raw) {#fun rleNms as rleNms_ { - withRLEs* `[RLE]', + withRLEs* `NonEmpty RLE', `Int', id `Ptr CUInt', `CDouble' } -> `()' #} -rleNms :: [RLE] -> Double -> IO [Bool] +rleNms :: NonEmpty RLE -> Double -> IO [Bool] rleNms dt thr = do let n = length dt allocaArray n $ \keep -> do @@ -234,20 +245,20 @@ bbNms :: BB -> Double -> IO [Bool] bbNms (BB dt) thr = do let n = SV.length dt - svUnsafeWith dt $ \pbb -> + svUnsafeWith dt $ \pbb -> allocaArray n $ \keep -> do bbNms_ (castPtr pbb) n keep (realToFrac thr) map (>0) <$> peekArray n keep {#fun rleToBbox as rleToBbox_ { - withRLEs* `[RLE]', + withRLEs* `NonEmpty RLE', `PtrBB', `Int' } -> `()' #} -rleToBbox :: [RLE] -> IO BB +rleToBbox :: NonEmpty RLE -> IO BB rleToBbox r = do let n = length r mbb <- SVM.new n @@ -264,7 +275,7 @@ } -> `()' #} -rleFrBbox :: BB -> Int -> Int -> IO [RLE] +rleFrBbox :: BB -> Int -> Int -> IO (NonEmpty RLE) rleFrBbox (BB bb) h w = do let n = SV.length bb makeRLEs n $ \prles -> svUnsafeWith bb $ \pbb -> do diff --git a/src/MXNet/Coco/Types.hs b/src/MXNet/Coco/Types.hs --- a/src/MXNet/Coco/Types.hs +++ b/src/MXNet/Coco/Types.hs @@ -2,14 +2,12 @@ {-# LANGUAGE TemplateHaskell #-} module MXNet.Coco.Types where -import Control.Applicative +import RIO hiding (Category) +import RIO.Char (ord) +import RIO.Time (LocalTime(..), TimeOfDay(..), Day, fromGregorianValid) import Data.Aeson import qualified Data.Attoparsec.Text as A -import Data.Time.Calendar (Day, fromGregorianValid) -import Data.Time (LocalTime(..), TimeOfDay(..)) import Data.Bits ((.&.)) -import Data.Char (ord) -import Data.Vector (Vector) import Control.Lens (makeLenses) import GHC.Generics (Generic) import Data.Store (Store) @@ -23,6 +21,7 @@ } deriving Generic instance Store Instance +instance NFData Instance instance FromJSON Instance where parseJSON = withObject "Instance" $ \v -> Instance @@ -42,6 +41,7 @@ } deriving Generic instance Store Info +instance NFData Info instance FromJSON Info where parseJSON = withObject "Info" $ \v -> Info @@ -59,6 +59,7 @@ } deriving Generic instance Store License +instance NFData License instance FromJSON License where parseJSON = withObject "License" $ \v -> License @@ -67,12 +68,12 @@ <*> v .: "url" data Image = Image { - _img_id :: !Int, - _img_width :: !Int, - _img_height :: !Int, - _img_file_name :: !String, - _img_license :: !Int, - _img_flickr_url :: !String, + _img_id :: !Int, + _img_width :: !Int, + _img_height :: !Int, + _img_file_name :: !String, + _img_license :: !Int, + _img_flickr_url :: !String, _img_coco_url :: !String, _img_date_captured :: !LocalTime } deriving (Generic, Show) @@ -82,6 +83,7 @@ instance Store TimeOfDay instance Store LocalTime instance Store Image +instance NFData Image instance FromJSON Image where parseJSON = withObject "Image" $ \v -> Image @@ -104,6 +106,7 @@ } deriving Generic instance Store Annotation +instance NFData Annotation instance FromJSON Annotation where parseJSON = withObject "Annotation" $ \v -> AnnObjectDetection @@ -118,11 +121,11 @@ deriving Generic instance Store Segmentation +instance NFData Segmentation instance FromJSON Segmentation where parseJSON value = (withObject "RLE" (\v -> SegRLE <$> v .: "counts" <*> v .: "size") value) <|> (withArray "Polygon" (\v -> SegPolygon <$> parseJSONList (Array v)) value) - data Category = CatObjectDetection { _odc_id :: Int, _odc_name :: String, @@ -130,6 +133,7 @@ } deriving Generic instance Store Category +instance NFData Category instance FromJSON Category where parseJSON = withObject "Category" $ \v -> CatObjectDetection @@ -140,6 +144,7 @@ newtype CocoDay = CocoDay Day deriving Generic instance Store CocoDay +instance NFData CocoDay instance FromJSON CocoDay where parseJSON = withText "Day" $ \t -> case A.parseOnly (day <* A.endOfInput) t of diff --git a/src/MXNet/NN/DataIter/Anchor.hs b/src/MXNet/NN/DataIter/Anchor.hs deleted file mode 100644 --- a/src/MXNet/NN/DataIter/Anchor.hs +++ /dev/null @@ -1,236 +0,0 @@ -{-# LANGUAGE TemplateHaskell #-} -module MXNet.NN.DataIter.Anchor where - -import qualified Data.IntSet as Set -import Control.Exception -import qualified Data.Vector as V -import qualified Data.Vector.Unboxed as UV -import qualified Data.Vector.Unboxed.Mutable as UVM -import Control.Lens (view, makeLenses) -import Control.Monad.Reader -import Data.Random (shuffleN, runRVar, StdRandom(..)) -import Data.Array.Repa (Array, DIM1, DIM2, D, U, (:.)(..), Z (..), All(..), (+^), fromListUnboxed) -import qualified Data.Array.Repa as Repa - --- import Debug.Trace - -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_allowed_border :: Int, - _conf_fg_num :: Int, - _conf_batch_num :: Int, - _conf_bg_overlap :: Float, - _conf_fg_overlap :: Float -} deriving Show -makeLenses ''Configuration - -anchors :: MonadReader Configuration m => - Int -> Int -> Int -> m (V.Vector (Anchor U)) -anchors stride width height = do - scales <- view conf_anchor_scales - ratios <- view conf_anchor_ratios - base <- baseAnchors stride - return $ V.fromList - [ Repa.computeS $ anch +^ offs - | offY <- grid height - , offX <- grid width - , anch <- base - , let offs = fromListUnboxed (Z :. 4) [offX, offY, offX, offY]] - where - grid size = map fromIntegral [0, stride .. size * stride-1] - -baseAnchors :: MonadReader Configuration m => - Int -> m ([Anchor U]) -baseAnchors size = do - scales <- view conf_anchor_scales - ratios <- view conf_anchor_ratios - return [makeBase s r | r <- ratios, s <- scales] - where - makeBase :: Int -> Float -> Anchor U - makeBase scale ratio = - let sizeF = fromIntegral 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) - -(#!) :: Array U DIM1 Float -> Int -> Float -(#!) = Repa.unsafeLinearIndex - -(%!) :: V.Vector a -> Int -> a -(%!) = (V.!) - -overlapMatrix :: Set.IntSet -> 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 DIM1 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) 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 - -- TODO case when gtBoxes is empty. - 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 j = UV.maxIndex $ Repa.toUnboxed $ Repa.computeS $ Repa.slice overlaps (Z :. i :. All) - let j = argMax overlaps 0 i - -- traceShowM $ ("GT -> ", j) - UVM.write labels j 1 - - -- 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 - -- let gts = UV.map (\i -> UV.maxIndex $ Repa.toUnboxed $ Repa.computeS $ Repa.slice overlaps (Z :. i :. All)) fgs - let gts = UV.map (argMax 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) labels - targetsRepa = Repa.fromUnboxed (Z:.numLabels:.4) (flattenT targets) - weightsRepa = Repa.fromUnboxed (Z:.numLabels:.4) (flattenT weights) - return (labelsRepa, targetsRepa, weightsRepa) - where - numGT = V.length gtBoxes - numLabels = V.length anBoxes - - argMax :: Array U DIM2 Float -> Int -> Int -> Int - argMax mat axis ind = - let series = case axis of - 0 -> Repa.slice mat $ Z :. ind :. All - 1 -> Repa.slice mat $ Z :. All :. ind - _ -> throw BadDimension - in UV.maxIndex $ Repa.toUnboxed $ Repa.computeS series - - 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.IntSet - 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) - - -- TODO: make it without any copy - flattenT :: UV.Vector (Float, Float, Float, Float) -> UV.Vector Float - flattenT = UV.concatMap (\(a,b,c,d) -> UV.fromList [a,b,c,d]) - -data AnchorError = BadDimension - deriving Show -instance Exception AnchorError diff --git a/src/MXNet/NN/DataIter/Coco.hs b/src/MXNet/NN/DataIter/Coco.hs deleted file mode 100644 --- a/src/MXNet/NN/DataIter/Coco.hs +++ /dev/null @@ -1,352 +0,0 @@ -{-# LANGUAGE TemplateHaskell #-} -{-# LANGUAGE DataKinds #-} -{-# LANGUAGE ScopedTypeVariables #-} -module MXNet.NN.DataIter.Coco ( - cocoImages, - cocoImagesWithAnchors, - cocoImagesWithAnchors', - Coco(..), - coco, -) where - -import Data.Maybe (catMaybes, fromMaybe) -import Data.List (unzip6) -import System.FilePath -import System.Directory -import GHC.Generics (Generic) -import qualified Data.ByteString as SBS -import qualified Data.Store as Store -import Control.Exception -import Data.Array.Repa (Array, DIM1, DIM3, D, U, (:.)(..), Z (..), Any(..), - fromListUnboxed, extent, backpermute, extend, (-^), (+^), (*^), (/^)) -import qualified Data.Array.Repa as Repa -import Data.Array.Repa.Repr.Unboxed (Unbox) -import qualified Data.Vector as V -import qualified Data.Vector.Unboxed as UV -import qualified Codec.Picture.Repa as RPJ -import Codec.Picture -import Codec.Picture.Extra -import qualified Data.Aeson as Aeson -import Control.Lens ((^.), (%~) , view, makeLenses, _1, _2) -import Data.Conduit -import qualified Data.Conduit.Combinators as C (yieldMany) -import qualified Data.Conduit.List as C -import Control.Monad.Reader -import qualified Data.IntMap.Strict as M -import Data.Maybe (fromJust) -import qualified Data.Random as RND (shuffleN, runRVar, StdRandom(..)) - -import MXNet.Base (NDArray(..), Fullfilled, ArgsHMap, ParameterList, Attr(..), (!), (!?), (.&), HMap(..), ArgOf(..), fromVector,zeros) -import MXNet.Base.Operators.NDArray (_Reshape) -import MXNet.NN.DataIter.Conduit -import qualified MXNet.NN.DataIter.Anchor as Anchor -import MXNet.Coco.Types - -data Coco = Coco FilePath String Instance - deriving Generic -instance Store.Store Coco - -raiseLeft :: Exception e => (a -> e) -> Either a b -> b -raiseLeft exc = either (throw . exc) id - -data FileNotFound = FileNotFound String String - deriving Show -instance Exception FileNotFound - -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 - return $ Coco base datasplit inst - -type ImageTensor = Array U DIM3 Float -type ImageInfo = Array U DIM1 Float -type GTBoxes = V.Vector (Array U DIM1 Float) - -data Configuration = Configuration { - _conf_short :: Int, - _conf_max_size :: Int, - _conf_mean :: (Float, Float, Float), - _conf_std :: (Float, Float, Float) -} -makeLenses ''Configuration - -cocoImages :: (MonadReader Configuration m, MonadIO m) => Coco -> Bool -> ConduitData m (ImageTensor, ImageInfo, GTBoxes) -cocoImages (Coco base datasplit inst) shuffle = ConduitData (Just 1) $ do - let all = inst ^. images - all_images <- if shuffle then - liftIO $ RND.runRVar (RND.shuffleN (length all) (V.toList all)) RND.StdRandom - else - return $ V.toList all - C.yieldMany all_images {-- .| C.iterM (liftIO . print) --} .| C.mapM loadImg .| C.catMaybes - where - -- dropAlpha tensor = - -- let Z :. _ :. w :. h = extent tensor - -- in fromFunction (Z :. (3 :: Int) :. w :. h) (tensor Repa.!) - loadImg img = do - short <- view conf_short - maxSize <- view conf_max_size - - let imgFilePath = base datasplit img ^. img_file_name - imgDyn <- raiseLeft (FileNotFound imgFilePath) <$> liftIO (readImage imgFilePath) - - let imgRGB = convertRGB8 imgDyn - imgH = fromIntegral $ imageHeight imgRGB - imgW = fromIntegral $ imageWidth imgRGB - - scale = calcScale imgW imgH short maxSize - imgH' = floor $ scale * imgH - imgW' = floor $ scale * imgW - imgInfo = fromListUnboxed (Z :. 3) [fromIntegral imgH', fromIntegral imgW', scale] - - imgResized = scaleBilinear imgW' imgH' imgRGB - imgRGBRepa = Repa.computeUnboxedS $ RPJ.imgData (RPJ.convertImage imgResized :: RPJ.Img RPJ.RGB) - - gt_boxes = get_gt_boxes scale img - - if V.null gt_boxes - then return Nothing - else do - let imgRGBRepa' = Repa.computeUnboxedS $ Repa.map fromIntegral imgRGBRepa - imgEval <- transform imgRGBRepa' - return $ Just (imgEval, imgInfo, gt_boxes) - - -- find a proper scale factor - calcScale imgW imgH short maxSize = - let imSizeMin = min imgH imgW - imSizeMax = max imgH imgW - imScale0 = fromIntegral short / imSizeMin :: Float - imScale1 = fromIntegral maxSize / imSizeMax :: Float - in if round (imScale0 * imSizeMax) > maxSize then imScale1 else imScale0 - - -- map each category from id to its index in the cocoClassNames. - catTabl = M.fromList $ V.toList $ V.map (\cat -> (cat ^. odc_id, fromJust $ V.elemIndex (cat ^. odc_name) cocoClassNames)) (inst ^. categories) - - -- get all the bbox and gt for the image - get_gt_boxes scale img = V.fromList $ catMaybes $ map makeGTBox $ V.toList imgAnns - where - imageId = img ^. img_id - width = img ^. img_width - height = img ^. img_height - imgAnns = V.filter (\ann -> ann ^. ann_image_id == imageId) (inst ^. annotations) - - 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) - - makeGTBox ann = - let (x0, y0, x1, y1) = cleanBBox (ann ^. ann_bbox) - classId = catTabl M.! (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 - - --- transform HWC -> CHW -transform :: MonadReader Configuration m => - Array U DIM3 Float -> m (Array U DIM3 Float) -transform img = do - mean <- view conf_mean - std <- view conf_std - let broadcast = Repa.computeUnboxedS . extend (Any :. height :. width) - mean' = broadcast $ fromTuple mean - std' = broadcast $ fromTuple std - chnFirst = backpermute newShape (\ (Z :. c :. h :. w) -> Z :. h :. w :. c) img - return $ Repa.computeUnboxedS $ (chnFirst -^ mean') /^ std' - where - (Z :. height :. width :. chn) = extent img - newShape = Z:. chn :. height :. width - --- transform CHW -> HWC -transformInv :: (Repa.Source r Float, MonadReader Configuration m) => - Array r DIM3 Float -> m (Array D DIM3 Float) -transformInv img = do - mean <- view conf_mean - std <- view conf_std - let broadcast = extend (Any :. height :. width) - mean' = broadcast $ fromTuple mean - std' = broadcast $ fromTuple std - addMean = img *^ std' +^ mean' - return $ backpermute newShape (\ (Z :. h :. w :. c) -> Z :. c :. h :. w) addMean - where - (Z :. chn :. height :. width) = extent img - newShape = Z :. height :. width :. chn - -fromTuple :: Unbox a => (a, a, a) -> Array U (Z :. Int) a -fromTuple (a, b, c) = fromListUnboxed (Z :. (3 :: Int)) [a,b,c] - -cocoClassNames = V.fromList [ - "__background__", -- always index 0 - "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"] - -type instance ParameterList "CocoImagesWithAnchors" = - '[ '("batch_size", 'AttrReq Int), - '("short_size", 'AttrReq Int), - '("long_size", 'AttrReq Int), - '("mean", 'AttrReq (Float, Float, Float)), - '("std", 'AttrReq (Float, Float, Float)), - '("feature_stride", 'AttrOpt Int), - '("anchor_scales", 'AttrOpt [Int]), - '("anchor_ratios", 'AttrOpt [Float]), - '("allowed_border", 'AttrOpt Int), - '("batch_rois", 'AttrOpt Int), - '("fg_fraction", 'AttrOpt Float), - '("fg_overlap", 'AttrOpt Float), - '("bg_overlap", 'AttrOpt Float), - '("shuffle", 'AttrOpt Bool)] - - -cocoImagesWithAnchors' :: (Fullfilled "CocoImagesWithAnchors" args, MonadIO m) => - ConduitData (ReaderT Configuration m) (ImageTensor, ImageInfo, GTBoxes) -> - ((Int, Int) -> IO (Int, Int)) -> - ArgsHMap "CocoImagesWithAnchors" args -> - ConduitData m ((NDArray Float, NDArray Float, NDArray Float), (NDArray Float, NDArray Float, NDArray Float)) -cocoImagesWithAnchors' (ConduitData _ images) extractFeatureShape args = ConduitData (Just batchSize) $ - morf images .| C.mapM (assignAnchors anchConf featureStride extractFeatureShape) .| C.chunksOf batchSize .| C.mapM toNDArray - where - batchSize = args ! #batch_size - batchRois = fromMaybe 256 $ args !? #batch_rois - featureStride = fromMaybe 16 $ args !? #feature_stride - anchConf = Anchor.Configuration { - Anchor._conf_anchor_scales = fromMaybe [8, 16, 32] $ args !? #anchor_scales, - Anchor._conf_anchor_ratios = fromMaybe [0.5, 1, 2] $ args !? #anchor_ratios, - Anchor._conf_allowed_border = fromMaybe 0 $ args !? #allowed_border, - Anchor._conf_fg_num = floor $ (fromMaybe 0.5 $ args !? #fg_fraction) * fromIntegral batchRois, - Anchor._conf_batch_num = batchRois, - Anchor._conf_fg_overlap = fromMaybe 0.7 $ args !? #fg_overlap, - Anchor._conf_bg_overlap = fromMaybe 0.3 $ args !? #bg_overlap - } - cocoConf = Configuration { - _conf_short = args ! #short_size, - _conf_max_size = args ! #long_size, - _conf_mean = args ! #mean, - _conf_std = args ! #std - } - morf = transPipe (flip runReaderT cocoConf) - -cocoImagesWithAnchors :: (Fullfilled "CocoImagesWithAnchors" args, MonadIO m) => - Coco -> ((Int, Int) -> IO (Int, Int)) -> ArgsHMap "CocoImagesWithAnchors" args -> - ConduitData m ((NDArray Float, NDArray Float, NDArray Float), (NDArray Float, NDArray Float, NDArray Float)) -cocoImagesWithAnchors cocoDef extractFeatureShape args = ConduitData (Just batchSize) $ - morf imgs .| C.mapM (assignAnchors anchConf featureStride extractFeatureShape) .| C.chunksOf batchSize .| C.mapM toNDArray - - where - ConduitData _ imgs = cocoImages cocoDef shuffle - cocoConf = Configuration { - _conf_short = args ! #short_size, - _conf_max_size = args ! #long_size, - _conf_mean = args ! #mean, - _conf_std = args ! #std - } - shuffle = fromMaybe True $ args !? #shuffle - batchSize = args ! #batch_size - batchRois = fromMaybe 256 $ args !? #batch_rois - featureStride = fromMaybe 16 $ args !? #feature_stride - anchConf = Anchor.Configuration { - Anchor._conf_anchor_scales = fromMaybe [8, 16, 32] $ args !? #anchor_scales, - Anchor._conf_anchor_ratios = fromMaybe [0.5, 1, 2] $ args !? #anchor_ratios, - Anchor._conf_allowed_border = fromMaybe 0 $ args !? #allowed_border, - Anchor._conf_fg_num = floor $ (fromMaybe 0.5 $ args !? #fg_fraction) * fromIntegral batchRois, - Anchor._conf_batch_num = batchRois, - Anchor._conf_fg_overlap = fromMaybe 0.7 $ args !? #fg_overlap, - Anchor._conf_bg_overlap = fromMaybe 0.3 $ args !? #bg_overlap - } - - morf = transPipe (flip runReaderT cocoConf) - -assignAnchors :: MonadIO m => Anchor.Configuration -> Int -> ((Int, Int) -> IO (Int, Int)) -> (ImageTensor, ImageInfo, GTBoxes) -> - m (ImageTensor, ImageInfo, GTBoxes, Repa.Array U DIM1 Float, Repa.Array U DIM3 Float, Repa.Array U DIM3 Float) -assignAnchors conf featureStride extractFeatureShape (img, info, gt) = do - let imHeight = floor $ info Anchor.#! 0 - imWidth = floor $ info Anchor.#! 1 - (featureWidth, featureHeight) <- liftIO $ extractFeatureShape (imWidth, imHeight) - anchors <- runReaderT (Anchor.anchors featureStride featureWidth featureHeight) conf - - (lbls, targets, weights) <- runReaderT (Anchor.assign gt imWidth imHeight anchors) conf - - -- reshape and transpose labls from (feat_h * feat_w * #anch, ) to (#anch, feat_h, feat_w) - -- reshape and transpose targets from (feat_h * feat_w * #anch, 4) to (#anch * 4, feat_h, feat_w) - -- reshape and transpose weights from (feat_h * feat_w * #anch, 4) to (#anch * 4, feat_h, feat_w) - let numAnch = length (conf ^. Anchor.conf_anchor_scales) * length (conf ^. Anchor.conf_anchor_ratios) - lbls <- return $ Repa.computeS $ - Repa.reshape (Z :. numAnch * featureHeight * featureWidth) $ - Repa.transpose $ - Repa.reshape (Z :. featureHeight * featureWidth :. numAnch) lbls - targets <- return $ Repa.computeS $ - Repa.reshape (Z :. numAnch * 4 :. featureHeight :. featureWidth) $ - Repa.transpose $ - Repa.reshape (Z :. featureHeight * featureWidth :. numAnch * 4) targets - weights <- return $ Repa.computeS $ - Repa.reshape (Z :. numAnch * 4 :. featureHeight :. featureWidth) $ - Repa.transpose $ - Repa.reshape (Z :. featureHeight * featureWidth :. numAnch * 4) weights - - -- let numAnch = length (anchConf ^. Anchor.conf_anchor_scales) * length (anchConf ^. Anchor.conf_anchor_ratios) - -- cvt1 (Z :. i :. h :. w) = Z :. ((h * featureWidth + w) * numAnch + i) - -- cvt2 (Z :. i :. h :. w) = let (m, n) = divMod i 4 - -- in Z :. ((h * featureWidth + w) * numAnch + m) :. n - -- lbls <- Repa.computeP $ Repa.reshape (Z :. numAnch * featureHeight :. featureWidth) $ - -- Repa.backpermute (Z :. numAnch :. featureHeight :. featureWidth) cvt1 lbls - -- targets <- Repa.computeP $ Repa.backpermute (Z :. numAnch * 4 :. featureHeight :. featureWidth) cvt2 targets - -- weights <- Repa.computeP $ Repa.backpermute (Z :. numAnch *4 :. featureHeight :. featureWidth) cvt2 weights - - return (img, info, gt, lbls, targets, weights) - --- toNDArray :: [((ImageTensor, ImageInfo, GTBoxes, Anchor.Labels, Anchor.Targets, Anchor.Weights))] -> --- IO ((NDArray Float, NDArray Float, NDArray Float), (NDArray Float, NDArray Float, NDArray Float)) -toNDArray dat = liftIO $ do - imagesC <- convertToMX images - infosC <- convertToMX infos - gtboxesC <- convertToMX [if V.null gt then emtpyGT else convertToRepa (V.toList gt) | gt <- gtboxes] - labelsC <- convertToMX labels - targetsC <- convertToMX targets - weightsC <- convertToMX weights - return ((imagesC, infosC, gtboxesC), (labelsC, targetsC, weightsC)) - where - (images, infos, gtboxes, labels, targets, weights) = unzip6 dat - emtpyGT = Repa.fromUnboxed (Z:.0:.5) UV.empty - - convert :: Repa.Shape sh => [Array U sh Float] -> ([Int], UV.Vector Float) - convert xs = assert (not (null xs)) $ (ext, vec) - where - vec = UV.concat $ map Repa.toUnboxed xs - sh0 = Repa.extent (head xs) - ext = length xs : reverse (Repa.listOfShape sh0) - - convertToMX :: Repa.Shape sh => [Array U sh Float] -> IO (NDArray Float) - convertToMX = uncurry fromVector . (_2 %~ UV.convert) . convert - - -- shape, at the type level, are sequence of Int, although we wnat to append - -- a dimension at the head, we add Int at the tail, they are the same. - convertToRepa :: Repa.Shape sh => [Array U sh Float] -> Array U (sh :. Int) Float - convertToRepa = uncurry Repa.fromUnboxed . (_1 %~ Repa.shapeOfList . reverse) . convert