{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE TemplateHaskell #-}

module RiskWeaver.Cmd.BDD where

import Control.Monad
import Control.Monad.Trans.Reader (ReaderT, ask, runReaderT)
import RiskWeaver.DSL.BDD qualified as BDD
import Data.ByteString qualified as BS
import Data.FileEmbed (embedFile)
import Data.List (sortBy)
import Data.Map qualified as Map
import Data.Text qualified as T
import Data.Vector (Vector)
import Data.Vector qualified as Vector
import RiskWeaver.Format.Coco
import RiskWeaver.Metric

import Options.Applicative
import System.Random
import Text.Printf
import RiskWeaver.Cmd.Core (RiskCommands(..))

cocoCategoryToClass :: CocoMap -> CategoryId -> BDD.Class
cocoCategoryToClass :: CocoMap -> CategoryId -> Class
cocoCategoryToClass CocoMap
coco CategoryId
categoryId =
  let cocoCategory :: CocoCategory
cocoCategory = (CocoMap -> Map CategoryId CocoCategory
cocoMapCocoCategory CocoMap
coco) Map CategoryId CocoCategory -> CategoryId -> CocoCategory
forall k a. Ord k => Map k a -> k -> a
Map.! CategoryId
categoryId
   in case Text -> String
T.unpack (CocoCategory -> Text
cocoCategoryName CocoCategory
cocoCategory) of
        String
"pedestrian" -> Class
BDD.Pedestrian
        String
"rider" -> Class
BDD.Rider
        String
"car" -> Class
BDD.Car
        String
"truck" -> Class
BDD.Truck
        String
"bus" -> Class
BDD.Bus
        String
"train" -> Class
BDD.Train
        String
"motorcycle" -> Class
BDD.Motorcycle
        String
"bicycle" -> Class
BDD.Bicycle
        String
_ -> Class
BDD.Background

cocoResultToVector :: CocoMap -> ImageId -> (Vector BDD.BoundingBoxGT, Vector BDD.BoundingBoxDT)
cocoResultToVector :: CocoMap -> ImageId -> (Vector BoundingBoxGT, Vector BoundingBoxDT)
cocoResultToVector CocoMap
coco ImageId
imageId = (Vector BoundingBoxGT
groundTruth, Vector BoundingBoxDT
detection)
  where
    groundTruth :: Vector BoundingBoxGT
groundTruth =
      [BoundingBoxGT] -> Vector BoundingBoxGT
forall a. [a] -> Vector a
Vector.fromList ([BoundingBoxGT] -> Vector BoundingBoxGT)
-> [BoundingBoxGT] -> Vector BoundingBoxGT
forall a b. (a -> b) -> a -> b
$
        [BoundingBoxGT]
-> ([CocoAnnotation] -> [BoundingBoxGT])
-> Maybe [CocoAnnotation]
-> [BoundingBoxGT]
forall b a. b -> (a -> b) -> Maybe a -> b
maybe
          []
          ( (CocoAnnotation -> BoundingBoxGT)
-> [CocoAnnotation] -> [BoundingBoxGT]
forall a b. (a -> b) -> [a] -> [b]
map
              ( \CocoAnnotation {Double
Int
Maybe Int
Maybe [[Double]]
CoCoBoundingBox
CategoryId
ImageId
cocoAnnotationId :: Int
cocoAnnotationImageId :: ImageId
cocoAnnotationCategory :: CategoryId
cocoAnnotationSegment :: Maybe [[Double]]
cocoAnnotationArea :: Double
cocoAnnotationBbox :: CoCoBoundingBox
cocoAnnotationIsCrowd :: Maybe Int
cocoAnnotationId :: CocoAnnotation -> Int
cocoAnnotationImageId :: CocoAnnotation -> ImageId
cocoAnnotationCategory :: CocoAnnotation -> CategoryId
cocoAnnotationSegment :: CocoAnnotation -> Maybe [[Double]]
cocoAnnotationArea :: CocoAnnotation -> Double
cocoAnnotationBbox :: CocoAnnotation -> CoCoBoundingBox
cocoAnnotationIsCrowd :: CocoAnnotation -> Maybe Int
..} ->
                  let CoCoBoundingBox (Double
cocox, Double
cocoy, Double
cocow, Double
cocoh) = CoCoBoundingBox
cocoAnnotationBbox
                   in BDD.BoundingBoxGT
                        { $sel:x:BoundingBoxGT :: Double
x = Double
cocox,
                          $sel:y:BoundingBoxGT :: Double
y = Double
cocoy,
                          $sel:w:BoundingBoxGT :: Double
w = Double
cocow,
                          $sel:h:BoundingBoxGT :: Double
h = Double
cocoh,
                          $sel:cls:BoundingBoxGT :: Class
cls = CocoMap -> CategoryId -> Class
cocoCategoryToClass CocoMap
coco CategoryId
cocoAnnotationCategory,
                          $sel:idx:BoundingBoxGT :: Int
idx = Int
cocoAnnotationId
                        }
              )
          )
          (ImageId -> Map ImageId [CocoAnnotation] -> Maybe [CocoAnnotation]
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup ImageId
imageId (CocoMap -> Map ImageId [CocoAnnotation]
cocoMapCocoAnnotation CocoMap
coco))
    detection :: Vector BoundingBoxDT
detection =
      [BoundingBoxDT] -> Vector BoundingBoxDT
forall a. [a] -> Vector a
Vector.fromList ([BoundingBoxDT] -> Vector BoundingBoxDT)
-> [BoundingBoxDT] -> Vector BoundingBoxDT
forall a b. (a -> b) -> a -> b
$
        [BoundingBoxDT]
-> ([CocoResult] -> [BoundingBoxDT])
-> Maybe [CocoResult]
-> [BoundingBoxDT]
forall b a. b -> (a -> b) -> Maybe a -> b
maybe
          []
          ( (CocoResult -> BoundingBoxDT) -> [CocoResult] -> [BoundingBoxDT]
forall a b. (a -> b) -> [a] -> [b]
map
              ( \CocoResult {CoCoBoundingBox
Score
CategoryId
ImageId
cocoResultImageId :: ImageId
cocoResultCategory :: CategoryId
cocoResultScore :: Score
cocoResultBbox :: CoCoBoundingBox
cocoResultImageId :: CocoResult -> ImageId
cocoResultCategory :: CocoResult -> CategoryId
cocoResultScore :: CocoResult -> Score
cocoResultBbox :: CocoResult -> CoCoBoundingBox
..} ->
                  let CoCoBoundingBox (Double
cocox, Double
cocoy, Double
cocow, Double
cocoh) = CoCoBoundingBox
cocoResultBbox
                   in BDD.BoundingBoxDT
                        { $sel:x:BoundingBoxDT :: Double
x = Double
cocox,
                          $sel:y:BoundingBoxDT :: Double
y = Double
cocoy,
                          $sel:w:BoundingBoxDT :: Double
w = Double
cocow,
                          $sel:h:BoundingBoxDT :: Double
h = Double
cocoh,
                          $sel:cls:BoundingBoxDT :: Class
cls = CocoMap -> CategoryId -> Class
cocoCategoryToClass CocoMap
coco CategoryId
cocoResultCategory,
                          $sel:score:BoundingBoxDT :: Double
score = Score -> Double
unScore Score
cocoResultScore,
                          $sel:idx:BoundingBoxDT :: Int
idx = ImageId -> Int
unImageId ImageId
cocoResultImageId
                        }
              )
          )
          (ImageId -> Map ImageId [CocoResult] -> Maybe [CocoResult]
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup ImageId
imageId (CocoMap -> Map ImageId [CocoResult]
cocoMapCocoResult CocoMap
coco))

runRisk ::
  CocoMap ->
  IO [(ImageId, Double)]
runRisk :: CocoMap -> IO [(ImageId, Double)]
runRisk CocoMap
cocoMap = do
  [ImageId]
-> (ImageId -> IO (ImageId, Double)) -> IO [(ImageId, Double)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM (CocoMap -> [ImageId]
cocoMapImageIds CocoMap
cocoMap) ((ImageId -> IO (ImageId, Double)) -> IO [(ImageId, Double)])
-> (ImageId -> IO (ImageId, Double)) -> IO [(ImageId, Double)]
forall a b. (a -> b) -> a -> b
$ \ImageId
imageId -> do
    let (Vector BoundingBoxGT
groundTruth, Vector BoundingBoxDT
detection) = CocoMap -> ImageId -> (Vector BoundingBoxGT, Vector BoundingBoxDT)
cocoResultToVector CocoMap
cocoMap ImageId
imageId
    let env :: Env BoundingBoxGT
env =
          BDD.MyEnv
            { $sel:envGroundTruth:MyEnv :: Vector BoundingBoxGT
envGroundTruth = Vector BoundingBoxGT
groundTruth,
              $sel:envDetection:MyEnv :: Vector BoundingBoxDT
envDetection = Vector BoundingBoxDT
detection,
              $sel:envConfidenceScoreThresh:MyEnv :: Double
envConfidenceScoreThresh = Double
0.4,
              $sel:envIoUThresh:MyEnv :: Double
envIoUThresh = Double
0.5
            }
    Double
risk <- (ReaderT (Env BoundingBoxGT) IO Double
 -> Env BoundingBoxGT -> IO Double)
-> Env BoundingBoxGT
-> ReaderT (Env BoundingBoxGT) IO Double
-> IO Double
forall a b c. (a -> b -> c) -> b -> a -> c
flip ReaderT (Env BoundingBoxGT) IO Double
-> Env BoundingBoxGT -> IO Double
forall r (m :: * -> *) a. ReaderT r m a -> r -> m a
runReaderT Env BoundingBoxGT
env (forall a (m :: * -> *).
(Fractional (Risk a), Num (Risk a), BoundingBox a, Monad m) =>
ReaderT (Env a) m (Risk a)
BDD.myRisk @BDD.BoundingBoxGT)
    (ImageId, Double) -> IO (ImageId, Double)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (ImageId
imageId, Double
risk)

showRisk :: Coco -> [CocoResult] -> Maybe Double -> Maybe Double -> Maybe ImageId -> IO ()
showRisk :: Coco
-> [CocoResult]
-> Maybe Double
-> Maybe Double
-> Maybe ImageId
-> IO ()
showRisk Coco
coco [CocoResult]
cocoResults Maybe Double
iouThreshold Maybe Double
scoreThresh Maybe ImageId
mImageId = do
  let cocoMap :: CocoMap
cocoMap =
        let cocoMap' :: CocoMap
cocoMap' = Coco -> [CocoResult] -> CocoMap
toCocoMap Coco
coco [CocoResult]
cocoResults
         in case Maybe ImageId
mImageId of
              Maybe ImageId
Nothing -> CocoMap
cocoMap'
              Just ImageId
imageId -> CocoMap
cocoMap' {cocoMapImageIds = [imageId]}
      iouThreshold' :: IOU
iouThreshold' = case Maybe Double
iouThreshold of
        Maybe Double
Nothing -> Double -> IOU
IOU Double
0.5
        Just Double
iouThreshold -> Double -> IOU
IOU Double
iouThreshold
      scoreThresh' :: Score
scoreThresh' = case Maybe Double
scoreThresh of
        Maybe Double
Nothing -> Double -> Score
Score Double
0.4
        Just Double
scoreThresh -> Double -> Score
Score Double
scoreThresh
  [(ImageId, Double)]
risks <- CocoMap -> IO [(ImageId, Double)]
runRisk CocoMap
cocoMap
  String -> IO ()
putStrLn (String -> IO ()) -> String -> IO ()
forall a b. (a -> b) -> a -> b
$ String -> String -> String -> String -> String
forall r. PrintfType r => String -> r
printf String
"%-12s %-12s %s" String
"#ImageId" String
"Filename" String
"Risk"
  let sortedRisks :: [(ImageId, Double)]
sortedRisks = ((ImageId, Double) -> (ImageId, Double) -> Ordering)
-> [(ImageId, Double)] -> [(ImageId, Double)]
forall a. (a -> a -> Ordering) -> [a] -> [a]
sortBy (\(ImageId
_, Double
risk1) (ImageId
_, Double
risk2) -> Double -> Double -> Ordering
forall a. Ord a => a -> a -> Ordering
compare Double
risk2 Double
risk1) [(ImageId, Double)]
risks
  [(ImageId, Double)] -> ((ImageId, Double) -> IO ()) -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [(ImageId, Double)]
sortedRisks (((ImageId, Double) -> IO ()) -> IO ())
-> ((ImageId, Double) -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \(ImageId
imageId, Double
risk) -> do
    let cocoImage :: CocoImage
cocoImage = (CocoMap -> Map ImageId CocoImage
cocoMapCocoImage CocoMap
cocoMap) Map ImageId CocoImage -> ImageId -> CocoImage
forall k a. Ord k => Map k a -> k -> a
Map.! ImageId
imageId
    String -> IO ()
putStrLn (String -> IO ()) -> String -> IO ()
forall a b. (a -> b) -> a -> b
$ String -> Int -> String -> Double -> String
forall r. PrintfType r => String -> r
printf String
"%-12d %-12s %.3f" (ImageId -> Int
unImageId ImageId
imageId) (Text -> String
T.unpack (CocoImage -> Text
cocoImageFileName CocoImage
cocoImage)) Double
risk

generateRiskWeightedDataset :: Coco -> [CocoResult] -> FilePath -> Maybe Double -> Maybe Double -> IO ()
generateRiskWeightedDataset :: Coco
-> [CocoResult] -> String -> Maybe Double -> Maybe Double -> IO ()
generateRiskWeightedDataset Coco
coco [CocoResult]
cocoResults String
cocoOutputFile Maybe Double
iouThreshold Maybe Double
scoreThresh = do
  let cocoMap :: CocoMap
cocoMap = Coco -> [CocoResult] -> CocoMap
toCocoMap Coco
coco [CocoResult]
cocoResults
      iouThreshold' :: IOU
iouThreshold' = case Maybe Double
iouThreshold of
        Maybe Double
Nothing -> Double -> IOU
IOU Double
0.5
        Just Double
iouThreshold -> Double -> IOU
IOU Double
iouThreshold
      scoreThresh' :: Score
scoreThresh' = case Maybe Double
scoreThresh of
        Maybe Double
Nothing -> Double -> Score
Score Double
0.4
        Just Double
scoreThresh -> Double -> Score
Score Double
scoreThresh
  [(ImageId, Double)]
risks <- CocoMap -> IO [(ImageId, Double)]
runRisk CocoMap
cocoMap
  let sumRisks :: Double
sumRisks = [Double] -> Double
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ([Double] -> Double) -> [Double] -> Double
forall a b. (a -> b) -> a -> b
$ ((ImageId, Double) -> Double) -> [(ImageId, Double)] -> [Double]
forall a b. (a -> b) -> [a] -> [b]
map (ImageId, Double) -> Double
forall a b. (a, b) -> b
snd [(ImageId, Double)]
risks
      probabilitis :: [Double]
probabilitis = ((ImageId, Double) -> Double) -> [(ImageId, Double)] -> [Double]
forall a b. (a -> b) -> [a] -> [b]
map (\(ImageId
_, Double
risk) -> Double
risk Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
sumRisks) [(ImageId, Double)]
risks
      accumulatedProbabilitis :: [Double]
accumulatedProbabilitis = (Double -> Double -> Double) -> Double -> [Double] -> [Double]
forall b a. (b -> a -> b) -> b -> [a] -> [b]
scanl Double -> Double -> Double
forall a. Num a => a -> a -> a
(+) Double
0 [Double]
probabilitis
      numDatasets :: Int
numDatasets = [ImageId] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ([ImageId] -> Int) -> [ImageId] -> Int
forall a b. (a -> b) -> a -> b
$ CocoMap -> [ImageId]
cocoMapImageIds CocoMap
cocoMap
      seed :: StdGen
seed = Int -> StdGen
mkStdGen Int
0

  -- Generate dataset by probability.
  -- The dataset's format is same as coco dataset.
  -- Accumurate probability
  -- Gen sorted list by accumulated probability with image id.
  -- Lottery by random number
  -- Get image id by lottery
  -- Generate random number between 0 and 1
  -- Find accumulated probability that is greater than random number
  -- Get image id by accumulated probability

  -- imageSets has accumulated probability and image id.
  -- It uses binary search to find image id by random number.
  let imageSets :: Vector (Double, ImageId)
      imageSets :: Vector (Double, ImageId)
imageSets = [(Double, ImageId)] -> Vector (Double, ImageId)
forall a. [a] -> Vector a
Vector.fromList ([(Double, ImageId)] -> Vector (Double, ImageId))
-> [(Double, ImageId)] -> Vector (Double, ImageId)
forall a b. (a -> b) -> a -> b
$ [Double] -> [ImageId] -> [(Double, ImageId)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Double]
accumulatedProbabilitis (CocoMap -> [ImageId]
cocoMapImageIds CocoMap
cocoMap)
      findImageIdFromImageSets :: Vector (Double, ImageId) -> Double -> ImageId
      findImageIdFromImageSets :: Vector (Double, ImageId) -> Double -> ImageId
findImageIdFromImageSets Vector (Double, ImageId)
imageSets Double
randomNum =
        let (Int
start, Int
end) = (Int
0, Vector (Double, ImageId) -> Int
forall a. Vector a -> Int
Vector.length Vector (Double, ImageId)
imageSets Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
            findImageIdFromImageSets' :: Int -> Int -> ImageId
            findImageIdFromImageSets' :: Int -> Int -> ImageId
findImageIdFromImageSets' Int
start Int
end =
              let mid :: Int
mid = (Int
start Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
end) Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
2
                  (Double
accumulatedProbability, ImageId
imageId) = Vector (Double, ImageId)
imageSets Vector (Double, ImageId) -> Int -> (Double, ImageId)
forall a. Vector a -> Int -> a
Vector.! Int
mid
               in if Int
start Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
end
                    then ImageId
imageId
                    else
                      if Double
accumulatedProbability Double -> Double -> Bool
forall a. Ord a => a -> a -> Bool
> Double
randomNum
                        then Int -> Int -> ImageId
findImageIdFromImageSets' Int
start Int
mid
                        else Int -> Int -> ImageId
findImageIdFromImageSets' (Int
mid Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Int
end
         in Int -> Int -> ImageId
findImageIdFromImageSets' Int
start Int
end
      lotteryN :: Int -> StdGen -> Int -> [ImageId]
      lotteryN :: Int -> StdGen -> Int -> [ImageId]
lotteryN Int
_ StdGen
_ Int
0 = []
      lotteryN Int
numDatasets StdGen
seed Int
n =
        let (Double
randNum, StdGen
seed') = (Double, Double) -> StdGen -> (Double, StdGen)
forall g. RandomGen g => (Double, Double) -> g -> (Double, g)
forall a g. (Random a, RandomGen g) => (a, a) -> g -> (a, g)
randomR (Double
0, Double
1) StdGen
seed
            imageId :: ImageId
imageId = Vector (Double, ImageId) -> Double -> ImageId
findImageIdFromImageSets Vector (Double, ImageId)
imageSets Double
randNum
         in ImageId
imageId ImageId -> [ImageId] -> [ImageId]
forall a. a -> [a] -> [a]
: Int -> StdGen -> Int -> [ImageId]
lotteryN Int
numDatasets StdGen
seed' (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
      imageIds :: [ImageId]
imageIds = Int -> StdGen -> Int -> [ImageId]
lotteryN Int
numDatasets StdGen
seed Int
numDatasets
      cocoImages' :: [CocoImage]
cocoImages' = (ImageId -> CocoImage) -> [ImageId] -> [CocoImage]
forall a b. (a -> b) -> [a] -> [b]
map (\ImageId
imageId -> (CocoMap -> Map ImageId CocoImage
cocoMapCocoImage CocoMap
cocoMap) Map ImageId CocoImage -> ImageId -> CocoImage
forall k a. Ord k => Map k a -> k -> a
Map.! ImageId
imageId) [ImageId]
imageIds
      newCoco :: Coco
newCoco =
        Coco
          { cocoImages :: [CocoImage]
cocoImages = [CocoImage]
cocoImages',
            cocoCategories :: [CocoCategory]
cocoCategories = Coco -> [CocoCategory]
cocoCategories Coco
coco,
            cocoAnnotations :: [CocoAnnotation]
cocoAnnotations = [],
            cocoLicenses :: Maybe [CocoLicense]
cocoLicenses = Coco -> Maybe [CocoLicense]
cocoLicenses Coco
coco,
            cocoInfo :: Maybe CocoInfo
cocoInfo = Coco -> Maybe CocoInfo
cocoInfo Coco
coco
          }
  String -> Coco -> IO ()
writeCoco String
cocoOutputFile Coco
newCoco


bddCommand :: RiskCommands
bddCommand :: RiskCommands
bddCommand = RiskCommands {
  showRisk :: Coco
-> [CocoResult]
-> Maybe Double
-> Maybe Double
-> Maybe ImageId
-> IO ()
showRisk = Coco
-> [CocoResult]
-> Maybe Double
-> Maybe Double
-> Maybe ImageId
-> IO ()
RiskWeaver.Cmd.BDD.showRisk,
  generateRiskWeightedDataset :: Coco
-> [CocoResult] -> String -> Maybe Double -> Maybe Double -> IO ()
generateRiskWeightedDataset = Coco
-> [CocoResult] -> String -> Maybe Double -> Maybe Double -> IO ()
RiskWeaver.Cmd.BDD.generateRiskWeightedDataset
}