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
          -- Use printf format to show score
          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
    -- 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 -> FilePath -> FilePath -> Bool -> IO ()
showImage Coco
coco FilePath
cocoFile FilePath
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 :: 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)