Merge branch 'master' of gitorious.org:hsv4l2/v4l2-examples
[hsv4l2:v4l2-examples.git] / v4l2-histogram / v4l2-histogram.hs
1 module Main (main) where
2
3 import Prelude hiding (map, sum)
4
5 import Control.Monad (when)
6 import Data.Word (Word8, Word32)
7 import Foreign (Ptr, nullPtr, allocaArray, peekElemOff, pokeElemOff)
8 import System.Exit (exitFailure, exitSuccess)
9 import System.IO (hPutStrLn, stderr)
10
11 import Graphics.UI.GLUT hiding (PixelFormat, histogram, imageHeight)
12 import Graphics.V4L2
13
14 pixel :: PixelFormat
15 pixel = PixelRGB24
16
17 main :: IO ()
18 main = do
19   initialWindowSize $= Size 640 480
20   initialDisplayMode $= [DoubleBuffered]
21   devname <- checkArgs . snd =<< getArgsAndInitialize
22   withDevice devname $ \d -> do
23     f <- setFormat d Capture . (\f->f{ imagePixelFormat = pixel }) =<< getFormat d Capture
24     checkFormat f
25     info $ "frame size: " ++ show (imageWidth f) ++ "x" ++ show (imageHeight f) ++ " pixels (" ++ show (imageSize f) ++ " bytes)"
26     let (_, texSize:_) = break (>= (imageWidth f `max` imageHeight f)) $ iterate (2*) 1
27     _ <- createWindow "v4l2-histogram"
28     depthFunc $= Nothing
29     texture Texture2D $= Enabled
30     [th, ti] <- genObjectNames 2
31     textureBinding Texture2D $= Just th
32     texImage2D Nothing NoProxy 0 RGBA' (TextureSize2D 256 256) 0 (PixelData RGBA UnsignedByte nullPtr)
33     textureFilter Texture2D $= ((Nearest, Nothing), Nearest)
34     textureWrapMode Texture2D S $= (Repeated, ClampToEdge)
35     textureWrapMode Texture2D T $= (Repeated, ClampToEdge)
36     textureBinding Texture2D $= Just ti
37     texImage2D Nothing NoProxy 0 RGBA' (TextureSize2D (fromIntegral texSize) (fromIntegral texSize)) 0 (PixelData RGBA UnsignedByte nullPtr)
38     textureFilter Texture2D $= ((Nearest, Nothing), Nearest)
39     textureWrapMode Texture2D S $= (Repeated, ClampToEdge)
40     textureWrapMode Texture2D T $= (Repeated, ClampToEdge)
41     matrixMode $= Modelview 0
42     loadIdentity
43     matrixMode $= Projection
44     loadIdentity
45     ortho2D 0 1 1 0
46     idleCallback $= Just (idle d f texSize ti th)
47     displayCallback $= display d f texSize ti th
48     keyboardMouseCallback $= Just (\_ _ _ _ -> exitSuccess)
49     mainLoop
50
51 idle :: Device -> ImageFormat -> Int -> TextureObject -> TextureObject -> IO ()
52 idle d f _ ti th = withFrame d f $ \p n -> do
53   if n == imageSize f
54     then allocaArray (256 * 3) $ \h -> allocaArray (256 * 256 * 4) $ \q -> do
55       histogram (imageWidth f * imageHeight f * 3) p h
56       expand (fromIntegral $ imageWidth f * imageHeight f) h q
57       textureBinding Texture2D $= Just ti
58       texSubImage2D Nothing 0 (TexturePosition2D 0 0) (TextureSize2D (fromIntegral $ imageWidth f) (fromIntegral $ imageHeight f)) (PixelData RGB UnsignedByte p)
59       textureBinding Texture2D $= Just th
60       texSubImage2D Nothing 0 (TexturePosition2D 0 0) (TextureSize2D 256 256) (PixelData RGBA UnsignedByte q)
61       postRedisplay Nothing      
62     else warn $ "incomplete frame (" ++ show n ++ " bytes, expected " ++ show (imageSize f) ++ " bytes)"
63
64 display:: Device -> ImageFormat -> Int -> TextureObject -> TextureObject -> IO ()
65 display _ f texSize ti th = do
66   textureBinding Texture2D $= Just ti
67   renderPrimitive Quads (u 0 0 >> u 0 1 >> u 1 1 >> u 1 0)
68   blend $= Enabled
69   blendFunc $= (SrcAlpha, OneMinusSrcAlpha)
70   textureBinding Texture2D $= Just th
71   renderPrimitive Quads (v 0 0 >> v 0 1 >> v 1 1 >> v 1 0)
72   blend $= Disabled
73   swapBuffers
74   where
75     u :: GLfloat -> GLfloat -> IO ()
76     u x y = texCoord (TexCoord2 (x * fromIntegral (imageWidth f) / fromIntegral texSize) (y * fromIntegral (imageHeight f) / fromIntegral texSize)) >> vertex (Vertex2 (1 - x) y)
77     v :: GLfloat -> GLfloat -> IO ()
78     v x y = texCoord (TexCoord2 x y) >> vertex (Vertex2 x y)
79
80 histogram :: Int -> Ptr Word8 -> Ptr Word32 -> IO ()
81 histogram m p q = c 0 >> h 0 >> h 1 >> h 2
82   where
83     c i | i >= 256 * 3 = return ()
84         | otherwise = do
85             pokeElemOff q i 0
86             c (i + 1)
87     h i0 = h' i0
88       where
89         h' i  | i >= m = return ()
90               | otherwise = do
91                   j <- peekElemOff p i
92                   let j' = fromIntegral j * 3 + i0
93                   t <- peekElemOff q j'
94                   pokeElemOff q j' (t + 1)
95                   h' (i + 3)
96
97 expand :: Float -> Ptr Word32 -> Ptr Word8 -> IO ()
98 expand m p q = e 0 >> a 0
99   where
100     e i | i >= 256 = return ()
101         | otherwise = e' 0 >> e' 1 >> e' 2 >> e (i + 1)
102       where
103         e' c = do
104           s <- peekElemOff p (3 * i + c)
105           let t | s == 0 = 255
106                 | otherwise = round . max 0 . min 255 $ 256 * log (m / fromIntegral s) / log 256
107           f t
108           where
109             f t = g 0
110               where
111                 g j | j == 256 = return ()
112                     | otherwise = do
113                         pokeElemOff q ((j * 256 + i) * 4 + c) $ if j > t then 255 else 0
114                         g (j + 1)
115     a i | i >= 256 * 256 * 4 = return ()
116         | otherwise = do
117             r <- peekElemOff q i
118             if r > 0 then pokeElemOff q (i + 3) 255 else do
119               g <- peekElemOff q (i + 1)
120               if g > 0 then pokeElemOff q (i + 3) 255 else do
121                 b <- peekElemOff q (i + 2)
122                 if b > 0 then pokeElemOff q (i + 3) 255 else do
123                   pokeElemOff q (i + 3) 0
124             a (i + 4)
125
126 checkFormat :: ImageFormat -> IO ()
127 checkFormat f = do
128   when (imagePixelFormat f /= pixel) $ err ("could not set pixel format " ++ show pixel)
129   when (imageBytesPerLine f /= imageWidth f * 3) $ err "cannot handle extra padding"
130   when (imageSize f /= imageBytesPerLine f * imageHeight f) $ err "cannot handle image size"
131
132 checkArgs :: [String] -> IO String
133 checkArgs [devname] = return devname
134 checkArgs _ = err $ "bad arguments; usage: v4l2-histogram /dev/video0"
135
136 err :: String -> IO a
137 err msg = (hPutStrLn stderr $ "**ERROR: [v4l2-histogram] " ++ msg) >> exitFailure
138
139 warn :: String -> IO ()
140 warn msg = hPutStrLn stderr $ "++ WARN: [v4l2-histogram] " ++ msg
141
142 info :: String -> IO ()
143 info msg = hPutStrLn stderr $ "   INFO: [v4l2-histogram] " ++ msg