From 70e3920556496e5fecb5fedddf1067b2522fcac7 Mon Sep 17 00:00:00 2001 From: Akshay Nair Date: Tue, 24 Dec 2024 18:51:17 +0530 Subject: Add setup for ocr with tesseract --- src/Chelleport/Context.hs | 8 ----- src/Chelleport/OCR.hs | 84 +++++++++++++++++++++++++++++++++++++++++++++++ src/Chelleport/Types.hs | 37 ++++++++++++++++++++- src/Chelleport/Utils.hs | 11 +++++++ 4 files changed, 131 insertions(+), 9 deletions(-) create mode 100644 src/Chelleport/OCR.hs (limited to 'src/Chelleport') diff --git a/src/Chelleport/Context.hs b/src/Chelleport/Context.hs index b23c1c2..30d8516 100644 --- a/src/Chelleport/Context.hs +++ b/src/Chelleport/Context.hs @@ -11,14 +11,6 @@ import SDL (($=)) import qualified SDL import qualified SDL.Font as TTF --- benchmark :: String -> IO a -> IO a --- benchmark msg m = do --- start <- systemNanoseconds <$> getSystemTime --- result <- m --- end <- systemNanoseconds <$> getSystemTime --- Debug.traceM $ msg ++ ": " ++ show (end - start) --- pure result - initializeContext :: IO DrawContext initializeContext = do -- Initialize SDL diff --git a/src/Chelleport/OCR.hs b/src/Chelleport/OCR.hs new file mode 100644 index 0000000..496b6a0 --- /dev/null +++ b/src/Chelleport/OCR.hs @@ -0,0 +1,84 @@ +module Chelleport.OCR (getWordsOnScreen) where + +import Chelleport.Types +import Control.Monad.IO.Class (MonadIO (liftIO)) +import Control.Monad.RWS (asks) +import qualified Data.ByteString as BS +import Foreign (Bits (shiftR), Ptr, Storable (peek, pokeByteOff), alloca, allocaBytes, peekArray, (.&.)) +import Foreign.C (CInt, CString, newCString) +import GHC.IO.Handle.FD (withFile) +import GHC.IO.IOMode (IOMode (WriteMode)) +import qualified Graphics.X11 as X11 +import qualified SDL +import System.Directory (removeFile) +import System.IO (hPutStrLn) +import System.IO.Temp (emptySystemTempFile) + +foreign import ccall unsafe "libchelleport.h findWordCoordinates" + c_findWordCoordinates :: CString -> Ptr CInt -> IO (Ptr OCRMatch) + +class (Monad m) => MonadOCR m where + getWordsOnScreen :: m [OCRMatch] + +instance (MonadIO m) => MonadOCR (AppM m) where + getWordsOnScreen = do + SDL.V2 width height <- asks ctxWindow >>= SDL.get . SDL.windowSize + SDL.V2 x y <- asks ctxWindow >>= SDL.getWindowAbsolutePosition + liftIO $ do + imgFilePath <- liftIO $ createTemporaryScreenshot (x, y) (width, height) + findWordCoordinates imgFilePath <* removeFile imgFilePath + +findWordCoordinates :: String -> IO [OCRMatch] +findWordCoordinates imgPath = alloca $ \sizePtr -> do + imgPathC <- newCString imgPath + arrayPtr <- c_findWordCoordinates imgPathC sizePtr + + size <- peek sizePtr + peekArray (fromIntegral size) arrayPtr + +createTemporaryScreenshot :: (CInt, CInt) -> (CInt, CInt) -> IO String +createTemporaryScreenshot offset size = do + tmpFilePath <- emptySystemTempFile "chelleport-screenshot.png" + screenshot tmpFilePath offset size + pure tmpFilePath + +screenshot :: String -> (CInt, CInt) -> (CInt, CInt) -> IO () +screenshot filename (offsetX, offsetY) (width, height) = do + dpy <- X11.openDisplay "" + root <- X11.rootWindow dpy (X11.defaultScreen dpy) + + image <- + X11.getImage + dpy + root + offsetX + offsetY + (fromIntegral width) + (fromIntegral height) + (fromIntegral X11.allPlanes_aux) + X11.zPixmap + + allocaBytes (fromIntegral $ width * height * 3) $ \ptr -> do + let getPixel :: CInt -> CInt -> IO () + getPixel x y = do + pixel <- X11.xGetPixel image x y + let r = pixel `shiftR` 16 .&. 0xFF + let g = pixel `shiftR` 8 .&. 0xFF + let b = pixel .&. 0xFF + pokeByteOff ptr (fromIntegral (y * width + x) * 3) r + pokeByteOff ptr (fromIntegral (y * width + x) * 3 + 1) g + pokeByteOff ptr (fromIntegral (y * width + x) * 3 + 2) b + + sequence_ [getPixel x y | y <- [0 .. height - 1], x <- [0 .. width - 1]] + rgbData <- BS.packCStringLen (ptr, fromIntegral $ width * height * 3) + savePPMFile filename (fromIntegral width) (fromIntegral height) rgbData + + X11.destroyImage image + X11.closeDisplay dpy + +savePPMFile :: FilePath -> Int -> Int -> BS.ByteString -> IO () +savePPMFile path width height rgbData = withFile path WriteMode $ \h -> do + hPutStrLn h "P6" + hPutStrLn h $ show width ++ " " ++ show height + hPutStrLn h "255" + BS.hPut h rgbData diff --git a/src/Chelleport/Types.hs b/src/Chelleport/Types.hs index 3c52909..cb580e0 100644 --- a/src/Chelleport/Types.hs +++ b/src/Chelleport/Types.hs @@ -1,7 +1,10 @@ module Chelleport.Types where import Control.Monad.Reader (MonadIO, MonadReader, ReaderT) +import Data.Vector.Storable (Storable) import Data.Word (Word8) +import Foreign (Ptr, Storable (alignment, peek, poke, sizeOf), castPtr, nullPtr, plusPtr) +import Foreign.C (CChar, CInt, peekCString) import qualified Graphics.X11 as X11 import qualified SDL import qualified SDL.Font as TTF @@ -14,13 +17,17 @@ type KeySequence = [Char] type KeyGrid = [[Cell]] +data Mode = ModeHints | ModeSearch + deriving (Show, Eq) + data State = State { stateGrid :: KeyGrid, stateKeySequence :: KeySequence, stateIsMatched :: Bool, stateIsShiftPressed :: Bool, stateIsDragging :: Bool, - stateRepetition :: Int + stateRepetition :: Int, + stateMode :: Mode } deriving (Show, Eq) @@ -51,3 +58,31 @@ data MouseButtonType = LeftClick | RightClick newtype AppM m a = AppM {runAppM :: ReaderT DrawContext m a} deriving (Functor, Applicative, Monad, MonadIO, MonadReader DrawContext) + +data OCRMatch = OCRMatch + { matchStartX :: !CInt, + matchStartY :: !CInt, + matchEndX :: !CInt, + matchEndY :: !CInt, + matchText :: !String + } + deriving (Show) + +instance Storable OCRMatch where + sizeOf _ = 4 * sizeOf (undefined :: CInt) + sizeOf (undefined :: Ptr CChar) + + -- TODO: Remove hardcoding later + alignment _ = 8 + + peek ptr = do + let cintSize = sizeOf (undefined :: CInt) + startX <- peek $ castPtr ptr + startY <- peek $ castPtr ptr `plusPtr` cintSize + endX <- peek $ castPtr ptr `plusPtr` (2 * cintSize) + endY <- peek $ castPtr ptr `plusPtr` (3 * cintSize) + text <- peek $ castPtr ptr `plusPtr` (4 * cintSize) + textStr <- if text == nullPtr then pure "" else peekCString text + pure $ OCRMatch startX startY endX endY textStr + + -- NOTE: Dont need poke + poke _ _ = undefined diff --git a/src/Chelleport/Utils.hs b/src/Chelleport/Utils.hs index 0e7dabc..c15a3e8 100644 --- a/src/Chelleport/Utils.hs +++ b/src/Chelleport/Utils.hs @@ -1,6 +1,9 @@ module Chelleport.Utils where +import Control.Monad.IO.Class (MonadIO (liftIO)) import Data.List (nub) +import Data.Time.Clock.System (SystemTime (systemNanoseconds), getSystemTime) +import qualified Debug.Trace as Debug import Foreign.C (CInt) intToCInt :: Int -> CInt @@ -24,3 +27,11 @@ isEmpty = null isNotEmpty :: [a] -> Bool isNotEmpty = not . isEmpty + +benchmark :: (MonadIO m) => String -> m a -> m a +benchmark msg m = do + start <- systemNanoseconds <$> liftIO getSystemTime + result <- m + end <- systemNanoseconds <$> liftIO getSystemTime + Debug.traceM $ msg ++ " (ms): " ++ show (fromIntegral (end - start) / 1_000_000.0 :: Double) + pure result -- cgit v1.3.1