{-# 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
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
}