module RiskWeaver.Display where
import Codec.Picture
import Control.Monad
import Data.Map qualified as Map
import Data.OSC1337 qualified as OSC
import Data.Sixel qualified as Sixel
import Data.Text qualified as T
import RiskWeaver.Draw
import RiskWeaver.Format.Coco
import System.Environment (lookupEnv)
import System.FilePath (takeBaseName, takeDirectory, (</>))
import Text.Printf
putImage :: Either FilePath (Image PixelRGB8) -> IO ()
putImage :: Either FilePath (Image PixelRGB8) -> IO ()
putImage Either FilePath (Image PixelRGB8)
image' = do
Maybe FilePath
termProgram <- FilePath -> IO (Maybe FilePath)
lookupEnv FilePath
"TERM_PROGRAM"
Image PixelRGB8
image <- case Either FilePath (Image PixelRGB8)
image' of
Left FilePath
imagePath -> do
Either FilePath DynamicImage
imageBin <- FilePath -> IO (Either FilePath DynamicImage)
readImage FilePath
imagePath
case Either FilePath DynamicImage
imageBin of
Left FilePath
err -> FilePath -> IO (Image PixelRGB8)
forall a. FilePath -> IO a
forall (m :: * -> *) a. MonadFail m => FilePath -> m a
fail (FilePath -> IO (Image PixelRGB8))
-> FilePath -> IO (Image PixelRGB8)
forall a b. (a -> b) -> a -> b
$ FilePath
"Image file " FilePath -> FilePath -> FilePath
forall a. [a] -> [a] -> [a]
++ FilePath
imagePath FilePath -> FilePath -> FilePath
forall a. [a] -> [a] -> [a]
++ FilePath
" can not be read."
Right DynamicImage
imageBin -> Image PixelRGB8 -> IO (Image PixelRGB8)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (DynamicImage -> Image PixelRGB8
convertRGB8 DynamicImage
imageBin)
Right Image PixelRGB8
image -> Image PixelRGB8 -> IO (Image PixelRGB8)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return Image PixelRGB8
image
case Maybe FilePath
termProgram of
Just FilePath
"iTerm.app" -> do
Image PixelRGB8 -> IO ()
forall a. ToOSC a => a -> IO ()
OSC.putOSC Image PixelRGB8
image
FilePath -> IO ()
putStrLn FilePath
""
Just FilePath
"vscode" -> do
Image PixelRGB8 -> IO ()
forall a. ToSixel a => a -> IO ()
Sixel.putSixel Image PixelRGB8
image
FilePath -> IO ()
putStrLn FilePath
""
Maybe FilePath
_ -> do
Image PixelRGB8 -> IO ()
forall a. ToSixel a => a -> IO ()
Sixel.putSixel Image PixelRGB8
image
FilePath -> IO ()
putStrLn FilePath
""
drawBoundingBox :: DynamicImage -> [CocoAnnotation] -> Map.Map CategoryId CocoCategory -> IO (Image PixelRGB8)
drawBoundingBox :: DynamicImage
-> [CocoAnnotation]
-> Map CategoryId CocoCategory
-> IO (Image PixelRGB8)
drawBoundingBox DynamicImage
imageBin [CocoAnnotation]
annotations Map CategoryId CocoCategory
categories = do
let imageRGB8 :: Image PixelRGB8
imageRGB8 = DynamicImage -> Image PixelRGB8
convertRGB8 DynamicImage
imageBin
[CocoAnnotation] -> (CocoAnnotation -> IO ()) -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [CocoAnnotation]
annotations ((CocoAnnotation -> IO ()) -> IO ())
-> (CocoAnnotation -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \CocoAnnotation
annotation -> do
let (CoCoBoundingBox (Double
bx, Double
by, Double
bw, Double
bh)) = CocoAnnotation -> CoCoBoundingBox
cocoAnnotationBbox CocoAnnotation
annotation
x :: Int
x = Double -> Int
forall b. Integral b => Double -> b
forall a b. (RealFrac a, Integral b) => a -> b
round Double
bx
y :: Int
y = Double -> Int
forall b. Integral b => Double -> b
forall a b. (RealFrac a, Integral b) => a -> b
round Double
by
width :: Int
width = Double -> Int
forall b. Integral b => Double -> b
forall a b. (RealFrac a, Integral b) => a -> b
round Double
bw
height :: Int
height = Double -> Int
forall b. Integral b => Double -> b
forall a b. (RealFrac a, Integral b) => a -> b
round Double
bh
Int
-> Int -> Int -> Int -> (Int, Int, Int) -> Image PixelRGB8 -> IO ()
drawRect Int
x Int
y (Int
x Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
width) (Int
y Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
height) (Int
255, Int
0, Int
0) Image PixelRGB8
imageRGB8
FilePath
-> Int
-> Int
-> (Int, Int, Int)
-> (Int, Int, Int)
-> Image PixelRGB8
-> IO ()
drawString (Text -> FilePath
T.unpack (CocoCategory -> Text
cocoCategoryName (Map CategoryId CocoCategory
categories Map CategoryId CocoCategory -> CategoryId -> CocoCategory
forall k a. Ord k => Map k a -> k -> a
Map.! CocoAnnotation -> CategoryId
cocoAnnotationCategory CocoAnnotation
annotation))) Int
x Int
y (Int
255, Int
0, Int
0) (Int
0, Int
0, Int
0) Image PixelRGB8
imageRGB8
Image PixelRGB8 -> IO (Image PixelRGB8)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return Image PixelRGB8
imageRGB8
drawDetectionBoundingBox :: DynamicImage -> [CocoResult] -> Map.Map CategoryId CocoCategory -> Maybe Double -> IO (Image PixelRGB8)
drawDetectionBoundingBox :: DynamicImage
-> [CocoResult]
-> Map CategoryId CocoCategory
-> Maybe Double
-> IO (Image PixelRGB8)
drawDetectionBoundingBox DynamicImage
imageBin [CocoResult]
annotations Map CategoryId CocoCategory
categories Maybe Double
scoreThreshold = do
let imageRGB8 :: Image PixelRGB8
imageRGB8 = DynamicImage -> Image PixelRGB8
convertRGB8 DynamicImage
imageBin
[CocoResult] -> (CocoResult -> IO ()) -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [CocoResult]
annotations ((CocoResult -> IO ()) -> IO ()) -> (CocoResult -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \CocoResult
annotation -> do
let (CoCoBoundingBox (Double
bx, Double
by, Double
bw, Double
bh)) = CocoResult -> CoCoBoundingBox
cocoResultBbox CocoResult
annotation
x :: Int
x = Double -> Int
forall b. Integral b => Double -> b
forall a b. (RealFrac a, Integral b) => a -> b
round Double
bx
y :: Int
y = Double -> Int
forall b. Integral b => Double -> b
forall a b. (RealFrac a, Integral b) => a -> b
round Double
by
width :: Int
width = Double -> Int
forall b. Integral b => Double -> b
forall a b. (RealFrac a, Integral b) => a -> b
round Double
bw
height :: Int
height = Double -> Int
forall b. Integral b => Double -> b
forall a b. (RealFrac a, Integral b) => a -> b
round Double
bh
draw :: IO ()
draw = do
Int
-> Int -> Int -> Int -> (Int, Int, Int) -> Image PixelRGB8 -> IO ()
drawRect Int
x Int
y (Int
x Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
width) (Int
y Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
height) (Int
255, Int
0, Int
0) Image PixelRGB8
imageRGB8
FilePath
-> Int
-> Int
-> (Int, Int, Int)
-> (Int, Int, Int)
-> Image PixelRGB8
-> IO ()
drawString (Text -> FilePath
T.unpack (CocoCategory -> Text
cocoCategoryName (Map CategoryId CocoCategory
categories Map CategoryId CocoCategory -> CategoryId -> CocoCategory
forall k a. Ord k => Map k a -> k -> a
Map.! CocoResult -> CategoryId
cocoResultCategory CocoResult
annotation))) Int
x Int
y (Int
255, Int
0, Int
0) (Int
0, Int
0, Int
0) Image PixelRGB8
imageRGB8
FilePath
-> Int
-> Int
-> (Int, Int, Int)
-> (Int, Int, Int)
-> Image PixelRGB8
-> IO ()
drawString (FilePath -> Double -> FilePath
forall r. PrintfType r => FilePath -> r
printf FilePath
"%.2f" (Score -> Double
unScore (Score -> Double) -> Score -> Double
forall a b. (a -> b) -> a -> b
$ CocoResult -> Score
cocoResultScore CocoResult
annotation)) Int
x (Int
y Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
10) (Int
255, Int
0, Int
0) (Int
0, Int
0, Int
0) Image PixelRGB8
imageRGB8
case Maybe Double
scoreThreshold of
Maybe Double
Nothing -> IO ()
draw
Just Double
scoreThreshold -> do
if CocoResult -> Score
cocoResultScore CocoResult
annotation Score -> Score -> Bool
forall a. Ord a => a -> a -> Bool
>= Double -> Score
Score Double
scoreThreshold
then IO ()
draw
else () -> IO ()
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ()
Image PixelRGB8 -> IO (Image PixelRGB8)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return Image PixelRGB8
imageRGB8
showImage :: Coco -> FilePath -> FilePath -> Bool -> IO ()
showImage :: Coco -> FilePath -> FilePath -> Bool -> IO ()
showImage Coco
coco FilePath
cocoFile FilePath
imageFile Bool
enableBoundingBox = do
let cocoFileNameWithoutExtension :: FilePath
cocoFileNameWithoutExtension = FilePath -> FilePath
takeBaseName FilePath
cocoFile
imageDir :: FilePath
imageDir = FilePath -> FilePath
takeDirectory (FilePath -> FilePath
takeDirectory FilePath
cocoFile) FilePath -> FilePath -> FilePath
</> FilePath
cocoFileNameWithoutExtension FilePath -> FilePath -> FilePath
</> FilePath
"images"
imagePath :: FilePath
imagePath = FilePath
imageDir FilePath -> FilePath -> FilePath
</> FilePath
imageFile
if Bool
enableBoundingBox
then do
let image' :: Maybe (CocoImage, [CocoAnnotation])
image' = Coco -> FilePath -> Maybe (CocoImage, [CocoAnnotation])
getCocoImageByFileName Coco
coco FilePath
imageFile
case Maybe (CocoImage, [CocoAnnotation])
image' of
Maybe (CocoImage, [CocoAnnotation])
Nothing -> FilePath -> IO ()
putStrLn (FilePath -> IO ()) -> FilePath -> IO ()
forall a b. (a -> b) -> a -> b
$ FilePath
"Image file " FilePath -> FilePath -> FilePath
forall a. [a] -> [a] -> [a]
++ FilePath
imageFile FilePath -> FilePath -> FilePath
forall a. [a] -> [a] -> [a]
++ FilePath
" is not found."
Just (CocoImage
image, [CocoAnnotation]
annotations) -> do
let categories :: Map CategoryId CocoCategory
categories = Coco -> Map CategoryId CocoCategory
toCategoryMap Coco
coco
Either FilePath DynamicImage
imageBin' <- FilePath -> IO (Either FilePath DynamicImage)
readImage FilePath
imagePath
case Either FilePath DynamicImage
imageBin' of
Left FilePath
err -> FilePath -> IO ()
putStrLn (FilePath -> IO ()) -> FilePath -> IO ()
forall a b. (a -> b) -> a -> b
$ FilePath
"Image file " FilePath -> FilePath -> FilePath
forall a. [a] -> [a] -> [a]
++ FilePath
imagePath FilePath -> FilePath -> FilePath
forall a. [a] -> [a] -> [a]
++ FilePath
" can not be read."
Right DynamicImage
imageBin -> do
Image PixelRGB8
imageRGB8 <- DynamicImage
-> [CocoAnnotation]
-> Map CategoryId CocoCategory
-> IO (Image PixelRGB8)
drawBoundingBox DynamicImage
imageBin [CocoAnnotation]
annotations Map CategoryId CocoCategory
categories
Either FilePath (Image PixelRGB8) -> IO ()
putImage (Image PixelRGB8 -> Either FilePath (Image PixelRGB8)
forall a b. b -> Either a b
Right Image PixelRGB8
imageRGB8)
else do
Either FilePath (Image PixelRGB8) -> IO ()
putImage (FilePath -> Either FilePath (Image PixelRGB8)
forall a b. a -> Either a b
Left FilePath
imagePath)
showDetectionImage :: Coco -> FilePath -> FilePath -> FilePath -> Maybe Double -> IO ()
showDetectionImage :: Coco -> FilePath -> FilePath -> FilePath -> Maybe Double -> IO ()
showDetectionImage Coco
coco FilePath
cocoFile FilePath
cocoResultFile FilePath
imageFile Maybe Double
scoreThreshold = do
let cocoFileNameWithoutExtension :: FilePath
cocoFileNameWithoutExtension = FilePath -> FilePath
takeBaseName FilePath
cocoFile
imageDir :: FilePath
imageDir = FilePath -> FilePath
takeDirectory (FilePath -> FilePath
takeDirectory FilePath
cocoFile) FilePath -> FilePath -> FilePath
</> FilePath
cocoFileNameWithoutExtension FilePath -> FilePath -> FilePath
</> FilePath
"images"
imagePath :: FilePath
imagePath = FilePath
imageDir FilePath -> FilePath -> FilePath
</> FilePath
imageFile
[CocoResult]
cocoResult <- FilePath -> IO [CocoResult]
readCocoResult FilePath
cocoResultFile
let image' :: Maybe (CocoImage, [CocoResult])
image' = Coco -> [CocoResult] -> FilePath -> Maybe (CocoImage, [CocoResult])
getCocoResultByFileName Coco
coco [CocoResult]
cocoResult FilePath
imageFile
case Maybe (CocoImage, [CocoResult])
image' of
Maybe (CocoImage, [CocoResult])
Nothing -> FilePath -> IO ()
putStrLn (FilePath -> IO ()) -> FilePath -> IO ()
forall a b. (a -> b) -> a -> b
$ FilePath
"Image file " FilePath -> FilePath -> FilePath
forall a. [a] -> [a] -> [a]
++ FilePath
imageFile FilePath -> FilePath -> FilePath
forall a. [a] -> [a] -> [a]
++ FilePath
" is not found."
Just (CocoImage
image, [CocoResult]
annotations) -> do
Either FilePath DynamicImage
imageBin' <- FilePath -> IO (Either FilePath DynamicImage)
readImage FilePath
imagePath
case Either FilePath DynamicImage
imageBin' of
Left FilePath
err -> FilePath -> IO ()
putStrLn (FilePath -> IO ()) -> FilePath -> IO ()
forall a b. (a -> b) -> a -> b
$ FilePath
"Image file " FilePath -> FilePath -> FilePath
forall a. [a] -> [a] -> [a]
++ FilePath
imagePath FilePath -> FilePath -> FilePath
forall a. [a] -> [a] -> [a]
++ FilePath
" can not be read."
Right DynamicImage
imageBin -> do
let categories :: Map CategoryId CocoCategory
categories = Coco -> Map CategoryId CocoCategory
toCategoryMap Coco
coco
Image PixelRGB8
imageRGB8 <- DynamicImage
-> [CocoResult]
-> Map CategoryId CocoCategory
-> Maybe Double
-> IO (Image PixelRGB8)
drawDetectionBoundingBox DynamicImage
imageBin ((CocoResult -> Bool) -> [CocoResult] -> [CocoResult]
forall a. (a -> Bool) -> [a] -> [a]
filter (\CocoResult
res -> CocoResult -> ImageId
cocoResultImageId CocoResult
res ImageId -> ImageId -> Bool
forall a. Eq a => a -> a -> Bool
== CocoImage -> ImageId
cocoImageId CocoImage
image) [CocoResult]
cocoResult) Map CategoryId CocoCategory
categories Maybe Double
scoreThreshold
Either FilePath (Image PixelRGB8) -> IO ()
putImage (Image PixelRGB8 -> Either FilePath (Image PixelRGB8)
forall a b. b -> Either a b
Right Image PixelRGB8
imageRGB8)