Display.hs

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 image' = do
  termProgram <- lookupEnv "TERM_PROGRAM"
  image <- case image' of
    Left imagePath -> do
      imageBin <- readImage imagePath
      case imageBin of
        Left err -> fail $ "Image file " ++ imagePath ++ " can not be read."
        Right imageBin -> return (convertRGB8 imageBin)
    Right image -> return image
  case termProgram of
    Just "iTerm.app" -> do
      OSC.putOSC image
      putStrLn ""
    Just "vscode" -> do
      Sixel.putSixel image
      putStrLn ""
    _ -> do
      Sixel.putSixel image
      putStrLn ""

drawBoundingBox :: DynamicImage -> [CocoAnnotation] -> Map.Map CategoryId CocoCategory -> IO (Image PixelRGB8)
drawBoundingBox imageBin annotations categories = do
  let imageRGB8 = convertRGB8 imageBin
  forM_ annotations $ \annotation -> do
    let (CoCoBoundingBox (bx, by, bw, bh)) = cocoAnnotationBbox annotation
        x = round bx
        y = round by
        width = round bw
        height = round bh
    drawRect x y (x + width) (y + height) (255, 0, 0) imageRGB8
    drawString (T.unpack (cocoCategoryName (categories Map.! cocoAnnotationCategory annotation))) x y (255, 0, 0) (0, 0, 0) imageRGB8
  return imageRGB8

drawDetectionBoundingBox :: DynamicImage -> [CocoResult] -> Map.Map CategoryId CocoCategory -> Maybe Double -> IO (Image PixelRGB8)
drawDetectionBoundingBox imageBin annotations categories scoreThreshold = do
  let imageRGB8 = convertRGB8 imageBin
  forM_ annotations $ \annotation -> do
    let (CoCoBoundingBox (bx, by, bw, bh)) = cocoResultBbox annotation
        x = round bx
        y = round by
        width = round bw
        height = round bh
        draw = do
          drawRect x y (x + width) (y + height) (255, 0, 0) imageRGB8
          drawString (T.unpack (cocoCategoryName (categories Map.! cocoResultCategory annotation))) x y (255, 0, 0) (0, 0, 0) imageRGB8
          -- Use printf format to show score
          drawString (printf "%.2f" (unScore $ cocoResultScore annotation)) x (y + 10) (255, 0, 0) (0, 0, 0) imageRGB8
    -- drawString (show $ cocoResultScore annotation)  x (y + 10) (255,0,0) (0,0,0) imageRGB8
    case scoreThreshold of
      Nothing -> draw
      Just scoreThreshold -> do
        if cocoResultScore annotation >= Score scoreThreshold
          then draw
          else return ()
  return imageRGB8

-- Show image by sixel
showImage :: Coco -> FilePath -> FilePath -> Bool -> IO ()
showImage coco cocoFile imageFile 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 = takeBaseName cocoFile
      imageDir = takeDirectory (takeDirectory cocoFile) </> cocoFileNameWithoutExtension </> "images"
      imagePath = imageDir </> imageFile
  if enableBoundingBox
    then do
      let image' = getCocoImageByFileName coco imageFile
      case image' of
        Nothing -> putStrLn $ "Image file " ++ imageFile ++ " is not found."
        Just (image, annotations) -> do
          let categories = toCategoryMap coco
          imageBin' <- readImage imagePath
          case imageBin' of
            Left err -> putStrLn $ "Image file " ++ imagePath ++ " can not be read."
            Right imageBin -> do
              imageRGB8 <- drawBoundingBox imageBin annotations categories
              putImage (Right imageRGB8)
    else do
      putImage (Left imagePath)

showDetectionImage :: Coco -> FilePath -> FilePath -> FilePath -> Maybe Double -> IO ()
showDetectionImage coco cocoFile cocoResultFile imageFile scoreThreshold = do
  let cocoFileNameWithoutExtension = takeBaseName cocoFile
      imageDir = takeDirectory (takeDirectory cocoFile) </> cocoFileNameWithoutExtension </> "images"
      imagePath = imageDir </> imageFile
  cocoResult <- readCocoResult cocoResultFile
  let image' = getCocoResultByFileName coco cocoResult imageFile
  case image' of
    Nothing -> putStrLn $ "Image file " ++ imageFile ++ " is not found."
    Just (image, annotations) -> do
      imageBin' <- readImage imagePath
      case imageBin' of
        Left err -> putStrLn $ "Image file " ++ imagePath ++ " can not be read."
        Right imageBin -> do
          let categories = toCategoryMap coco
          imageRGB8 <- drawDetectionBoundingBox imageBin (filter (\res -> cocoResultImageId res == cocoImageId image) cocoResult) categories scoreThreshold
          putImage (Right imageRGB8)