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 String (Image PixelRGB8) -> IO ()
putImage Either String (Image PixelRGB8)
image' = do
Maybe String
termProgram <- String -> IO (Maybe String)
lookupEnv String
"TERM_PROGRAM"
Image PixelRGB8
image <- case Either String (Image PixelRGB8)
image' of
Left String
imagePath -> do
Either String DynamicImage
imageBin <- String -> IO (Either String DynamicImage)
readImage String
imagePath
case Either String DynamicImage
imageBin of
Left String
err -> String -> IO (Image PixelRGB8)
forall a. String -> IO a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail (String -> IO (Image PixelRGB8)) -> String -> IO (Image PixelRGB8)
forall a b. (a -> b) -> a -> b
$ String
"Image file " String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
imagePath String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
" can not be read. : " String -> String -> String
forall a. [a] -> [a] -> [a]
++ String -> String
forall a. Show a => a -> String
show String
err
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 String
termProgram of
Just String
"iTerm.app" -> do
Image PixelRGB8 -> IO ()
forall a. ToOSC a => a -> IO ()
OSC.putOSC Image PixelRGB8
image
String -> IO ()
putStrLn String
""
Just String
"vscode" -> do
Image PixelRGB8 -> IO ()
forall a. ToSixel a => a -> IO ()
Sixel.putSixel Image PixelRGB8
image
String -> IO ()
putStrLn String
""
Maybe String
_ -> do
Image PixelRGB8 -> IO ()
forall a. ToSixel a => a -> IO ()
Sixel.putSixel Image PixelRGB8
image
String -> IO ()
putStrLn String
""
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
String
-> Int
-> Int
-> (Int, Int, Int)
-> (Int, Int, Int)
-> Image PixelRGB8
-> IO ()
drawString (Text -> String
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 ::
(Show a) =>
DynamicImage ->
[CocoResult] ->
[a] ->
Map.Map CategoryId CocoCategory ->
Maybe Double ->
Maybe (Image PixelRGB8 -> a -> IO (Image PixelRGB8)) ->
IO (Image PixelRGB8)
drawDetectionBoundingBox :: forall a.
Show a =>
DynamicImage
-> [CocoResult]
-> [a]
-> Map CategoryId CocoCategory
-> Maybe Double
-> Maybe (Image PixelRGB8 -> a -> IO (Image PixelRGB8))
-> IO (Image PixelRGB8)
drawDetectionBoundingBox DynamicImage
imageBin [CocoResult]
annotations [a]
properties Map CategoryId CocoCategory
categories Maybe Double
scoreThreshold Maybe (Image PixelRGB8 -> a -> IO (Image PixelRGB8))
overlay = do
let imageRGB8 :: Image PixelRGB8
imageRGB8 = DynamicImage -> Image PixelRGB8
convertRGB8 DynamicImage
imageBin
zipedAnnotations :: [(CocoResult, Maybe a)]
zipedAnnotations = [CocoResult] -> [Maybe a] -> [(CocoResult, Maybe a)]
forall a b. [a] -> [b] -> [(a, b)]
zip [CocoResult]
annotations ([Maybe a] -> [(CocoResult, Maybe a)])
-> [Maybe a] -> [(CocoResult, Maybe a)]
forall a b. (a -> b) -> a -> b
$
case [a]
properties of
[] -> Maybe a -> [Maybe a]
forall a. a -> [a]
repeat Maybe a
forall a. Maybe a
Nothing
[a]
_ -> (a -> Maybe a) -> [a] -> [Maybe a]
forall a b. (a -> b) -> [a] -> [b]
map a -> Maybe a
forall a. a -> Maybe a
Just [a]
properties
[(CocoResult, Maybe a)]
-> ((CocoResult, Maybe a) -> IO ()) -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [(CocoResult, Maybe a)]
zipedAnnotations (((CocoResult, Maybe a) -> IO ()) -> IO ())
-> ((CocoResult, Maybe a) -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \(CocoResult
annotation, Maybe a
property) -> 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
String
-> Int
-> Int
-> (Int, Int, Int)
-> (Int, Int, Int)
-> Image PixelRGB8
-> IO ()
drawString (Text -> String
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
String
-> Int
-> Int
-> (Int, Int, Int)
-> (Int, Int, Int)
-> Image PixelRGB8
-> IO ()
drawString (String -> Double -> String
forall r. PrintfType r => String -> r
printf String
"%.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 a
property, Maybe (Image PixelRGB8 -> a -> IO (Image PixelRGB8))
overlay) of
(Just a
property', Just Image PixelRGB8 -> a -> IO (Image PixelRGB8)
overlay') -> do
Image PixelRGB8
_ <- Image PixelRGB8 -> a -> IO (Image PixelRGB8)
overlay' Image PixelRGB8
imageRGB8 a
property'
() -> IO ()
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ()
(Maybe a, Maybe (Image PixelRGB8 -> a -> IO (Image PixelRGB8)))
_ -> () -> IO ()
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ()
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 -> String -> String -> Bool -> IO ()
showImage Coco
coco String
cocoFile String
imageFile Bool
enableBoundingBox = do
let cocoFileNameWithoutExtension :: String
cocoFileNameWithoutExtension = String -> String
takeBaseName String
cocoFile
imageDir :: String
imageDir = String -> String
takeDirectory (String -> String
takeDirectory String
cocoFile) String -> String -> String
</> String
cocoFileNameWithoutExtension String -> String -> String
</> String
"images"
imagePath :: String
imagePath = String
imageDir String -> String -> String
</> String
imageFile
if Bool
enableBoundingBox
then do
let image' :: Maybe (CocoImage, [CocoAnnotation])
image' = Coco -> String -> Maybe (CocoImage, [CocoAnnotation])
getCocoImageByFileName Coco
coco String
imageFile
case Maybe (CocoImage, [CocoAnnotation])
image' of
Maybe (CocoImage, [CocoAnnotation])
Nothing -> String -> IO ()
putStrLn (String -> IO ()) -> String -> IO ()
forall a b. (a -> b) -> a -> b
$ String
"Image file " String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
imageFile String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
" is not found."
Just (CocoImage
_, [CocoAnnotation]
annotations) -> do
let categories :: Map CategoryId CocoCategory
categories = Coco -> Map CategoryId CocoCategory
toCategoryMap Coco
coco
Either String DynamicImage
imageBin' <- String -> IO (Either String DynamicImage)
readImage String
imagePath
case Either String DynamicImage
imageBin' of
Left String
err -> String -> IO ()
putStrLn (String -> IO ()) -> String -> IO ()
forall a b. (a -> b) -> a -> b
$ String
"Image file " String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
imagePath String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
" can not be read. : " String -> String -> String
forall a. [a] -> [a] -> [a]
++ String -> String
forall a. Show a => a -> String
show String
err
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 String (Image PixelRGB8) -> IO ()
putImage (Image PixelRGB8 -> Either String (Image PixelRGB8)
forall a b. b -> Either a b
Right Image PixelRGB8
imageRGB8)
else do
Either String (Image PixelRGB8) -> IO ()
putImage (String -> Either String (Image PixelRGB8)
forall a b. a -> Either a b
Left String
imagePath)
showDetectionImage :: (Show a) => CocoMap -> FilePath -> Maybe Double -> [a] -> Maybe (Image PixelRGB8 -> a -> IO (Image PixelRGB8)) -> IO ()
showDetectionImage :: forall a.
Show a =>
CocoMap
-> String
-> Maybe Double
-> [a]
-> Maybe (Image PixelRGB8 -> a -> IO (Image PixelRGB8))
-> IO ()
showDetectionImage CocoMap
cocoMap String
imageFile Maybe Double
scoreThreshold [a]
properties Maybe (Image PixelRGB8 -> a -> IO (Image PixelRGB8))
overlay = do
let imagePath :: String
imagePath = CocoMap -> String
getImageDir CocoMap
cocoMap String -> String -> String
</> String
imageFile
let image' :: Maybe (CocoImage, [CocoResult])
image' = CocoMap -> String -> Maybe (CocoImage, [CocoResult])
forall a.
CocoMapable a =>
CocoMap -> a -> Maybe (CocoImage, [CocoResult])
getCocoResult CocoMap
cocoMap String
imageFile
case Maybe (CocoImage, [CocoResult])
image' of
Maybe (CocoImage, [CocoResult])
Nothing -> String -> IO ()
putStrLn (String -> IO ()) -> String -> IO ()
forall a b. (a -> b) -> a -> b
$ String
"Image file " String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
imageFile String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
" is not found."
Just (CocoImage
_, [CocoResult]
annotations) -> do
Either String DynamicImage
imageBin' <- String -> IO (Either String DynamicImage)
readImage String
imagePath
case Either String DynamicImage
imageBin' of
Left String
err -> String -> IO ()
putStrLn (String -> IO ()) -> String -> IO ()
forall a b. (a -> b) -> a -> b
$ String
"Image file " String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
imagePath String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
" can not be read. : " String -> String -> String
forall a. [a] -> [a] -> [a]
++ String -> String
forall a. Show a => a -> String
show String
err
Right DynamicImage
imageBin -> do
DynamicImage
-> [CocoResult]
-> [a]
-> Map CategoryId CocoCategory
-> Maybe Double
-> Maybe (Image PixelRGB8 -> a -> IO (Image PixelRGB8))
-> IO (Image PixelRGB8)
forall a.
Show a =>
DynamicImage
-> [CocoResult]
-> [a]
-> Map CategoryId CocoCategory
-> Maybe Double
-> Maybe (Image PixelRGB8 -> a -> IO (Image PixelRGB8))
-> IO (Image PixelRGB8)
drawDetectionBoundingBox DynamicImage
imageBin [CocoResult]
annotations [a]
properties (CocoMap -> Map CategoryId CocoCategory
cocoMapCocoCategory CocoMap
cocoMap) Maybe Double
scoreThreshold Maybe (Image PixelRGB8 -> a -> IO (Image PixelRGB8))
overlay
IO (Image PixelRGB8) -> (Image PixelRGB8 -> IO ()) -> IO ()
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Either String (Image PixelRGB8) -> IO ()
putImage (Either String (Image PixelRGB8) -> IO ())
-> (Image PixelRGB8 -> Either String (Image PixelRGB8))
-> Image PixelRGB8
-> IO ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Image PixelRGB8 -> Either String (Image PixelRGB8)
forall a b. b -> Either a b
Right