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) =>
  -- | Image
  DynamicImage ->
  -- | A list of Coco result
  [CocoResult] ->
  -- | A list of object property
  [a] ->
  -- | A map of category
  Map.Map CategoryId CocoCategory ->
  -- | Score threshold
  Maybe Double ->
  -- | Overlay function to draw object property
  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
          -- Use printf format to show score
          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 ()
    -- drawString (show $ cocoResultScore annotation)  x (y + 10) (255,0,0) (0,0,0) 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

-- Show image by sixel
showImage :: Coco -> FilePath -> FilePath -> Bool -> IO ()
showImage :: Coco -> String -> String -> Bool -> IO ()
showImage Coco
coco String
cocoFile String
imageFile Bool
enableBoundingBox = do
  -- Get a diretory of image file from cocoFile's filename.
  -- cocoFile's filename is the same as image directory.
  -- For example, cocoFile is annotations/test.json, then image directory is test/images, and lable directory is test/labels.
  -- Get a parent parent directory(grand parent directory) of cocoFile's filename, and add a directory of images
  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