potatonization.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": true,
    "deletable": true,
    "editable": true
   },
   "outputs": [],
   "source": [
    "from ctypes import cast, c_void_p, c_byte, ARRAY, POINTER\n",
    "import time\n",
    "import threading\n",
    "\n",
    "import numpy as np\n",
    "import cv2\n",
    "\n",
    "from pynq import Overlay\n",
    "from pynq.drivers import HDMI\n",
    "from pynq.drivers import video\n",
    "\n",
    "import threading"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": true,
    "deletable": true,
    "editable": true
   },
   "outputs": [],
   "source": [
    "POTETOFILES = [\n",
    "    'img/potato_256.png',\n",
    "    'img/potato_128.png',\n",
    "    'img/potato_64.png'\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "collapsed": true,
    "deletable": true,
    "editable": true
   },
   "outputs": [],
   "source": [
    "Overlay('base.bit').download()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "collapsed": true,
    "deletable": true,
    "editable": true
   },
   "outputs": [],
   "source": [
    "class HDMIFilter(object):\n",
    "    MAX_SHAPE = (1080, 1920, 3)\n",
    "    BUFFER_SIZE = MAX_SHAPE[0]*MAX_SHAPE[1]*MAX_SHAPE[2]\n",
    "    TIMEOUT = 1\n",
    "    \n",
    "    def __init__(self, filter_fn, video_mode=video.VMODE_1280x720):\n",
    "        self.video_mode = video_mode\n",
    "        self.hdmi_out = HDMI('out', video_mode=video_mode)\n",
    "        try:\n",
    "            self.hdmi_in = HDMI('in', init_timeout=10, frame_list=self.hdmi_out.frame_list)\n",
    "        except Exception as e:\n",
    "            print(e)\n",
    "            self.hdmi_in = None\n",
    "        \n",
    "        self.filter_fn = filter_fn\n",
    "        self.frames = []\n",
    "        for i in range(3):\n",
    "            addr = self.hdmi_out.frame_addr(i)\n",
    "            void_p = c_void_p(addr)\n",
    "            buffer = cast(void_p, POINTER(ARRAY(c_byte, len=self.BUFFER_SIZE))).contents\n",
    "            self.frames.append(\n",
    "                np.frombuffer(buffer, dtype=np.uint8).reshape(self.MAX_SHAPE)[\n",
    "                    :self.hdmi_out.frame_height(),\n",
    "                    :self.hdmi_out.frame_width()\n",
    "                ]\n",
    "            )\n",
    "    \n",
    "    def start(self):\n",
    "        if self.hdmi_in is not None:\n",
    "            self.hdmi_in.start()\n",
    "        self.hdmi_out.start()\n",
    "        return self\n",
    "    \n",
    "    def stop(self):\n",
    "        if self.hdmi_in is not None:\n",
    "            self.hdmi_in.stop()\n",
    "        self.hdmi_out.stop()\n",
    "    \n",
    "    def close(self):\n",
    "        self.stop()\n",
    "    \n",
    "    def run(self, timeout=None):\n",
    "        self.start()\n",
    "        start = time.time()\n",
    "        while True:\n",
    "            if timeout is not None:\n",
    "                now = time.time()\n",
    "                if start + timeout <= now:\n",
    "                    return self.TIMEOUT\n",
    "            if self.hdmi_in is not None:\n",
    "                index = self.hdmi_in.frame_index()\n",
    "                self.hdmi_in.frame_index_next()\n",
    "                time.sleep(1./120)\n",
    "            else:\n",
    "                index = (self.hdmi_out.frame_index() + 1)%3\n",
    "            frame = self.frames[index]\n",
    "            self.filter_fn(frame)\n",
    "            time.sleep(1./120)\n",
    "            self.hdmi_out.frame_index(index)\n",
    "        self.stop()\n",
    "        \n",
    "    def frame(self, idx=None):\n",
    "        idx = idx if idx else self.hdmi_in.frame_index()\n",
    "        return self.frames[idx]\n",
    "        \n",
    "    def __enter__(self):\n",
    "        self.start()\n",
    "        return self\n",
    "    \n",
    "    def __exit__(self, exc_type, exc_value, traceback):\n",
    "        self.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "collapsed": true,
    "deletable": true,
    "editable": true
   },
   "outputs": [],
   "source": [
    "class PotatoOverWriter(object):\n",
    "    def __init__(self, potato_files, faces):\n",
    "        imgs = []\n",
    "        for potato_file in potato_files:\n",
    "            img = cv2.imread(potato_file, cv2.IMREAD_UNCHANGED)\n",
    "            size = min(img.shape[:2])\n",
    "            imgs.append((size, img.shape, img))\n",
    "        self.imgs = sorted(imgs, key=lambda x: x[0])\n",
    "        self.faces = faces\n",
    "        self.cnt = 0\n",
    "    \n",
    "    def __call__(self, img):\n",
    "        for x, y, w, h in self.faces:\n",
    "            face_size = max(w, h)\n",
    "            for potato_size, potato_shape, potato_img in self.imgs:\n",
    "                if face_size <= potato_size:\n",
    "                    break\n",
    "            pos = (\n",
    "                max(0, y + (h - potato_shape[0])//2),\n",
    "                max(0, x + (w - potato_shape[1])//2)\n",
    "            )\n",
    "            self._overlay(img, potato_img, pos)\n",
    "    \n",
    "    def _overlay(self, img1, img2, pos):\n",
    "        window = img1[pos[0]:, pos[1]:][:img2.shape[0], :img2.shape[1]]\n",
    "        mask = (img2[:, :, 3] > 127).astype(np.uint8)[:window.shape[0], :window.shape[1]]\n",
    "        if min(window.shape) > 0:\n",
    "            for ch in range(3):\n",
    "                window[:, :, ch] *= (1 - mask)\n",
    "                window[:, :, ch] += mask*img2[:window.shape[0], :window.shape[1], ch]\n",
    "        return img1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "collapsed": false,
    "deletable": true,
    "editable": true
   },
   "outputs": [],
   "source": [
    "class FaceDetectionThread(threading.Thread):\n",
    "    def __init__(self, faces, hdmi_filter, cascade='haarcascade_frontalface_default.xml'):\n",
    "        super(self.__class__, self).__init__()\n",
    "        self.cascade = cascade\n",
    "        self._stop_event = threading.Event()\n",
    "        self._face_cascade = cv2.CascadeClassifier(cascade)\n",
    "        self.faces = faces\n",
    "        self.hdmi_filter = hdmi_filter\n",
    "\n",
    "    def run(self):\n",
    "        while not self.stopped():\n",
    "            frame = self.hdmi_filter.frame()\n",
    "            gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)\n",
    "            faces = self._face_cascade.detectMultiScale(gray, 1.3, 5)\n",
    "            self.faces.clear()\n",
    "            self.faces.extend(faces)\n",
    "        \n",
    "    def stop(self):\n",
    "        self._stop_event.set()\n",
    "\n",
    "    def stopped(self):\n",
    "        return self._stop_event.is_set()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "collapsed": false,
    "deletable": true,
    "editable": true
   },
   "outputs": [],
   "source": [
    "faces = []\n",
    "\n",
    "with HDMIFilter(PotatoOverWriter(POTETOFILES, faces)) as f_:\n",
    "    face_detector = FaceDetectionThread(faces, f_)\n",
    "    face_detector.start()\n",
    "    f_.run(60)\n",
    "    face_detector.stop()\n",
    "    face_detector.join()\n",
    "    hdmi_out = f_.hdmi_out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "f = np.frombuffer(bytearray(hdmi_out.frame_raw()), dtype=np.uint8).reshape((1080, 1920, 3))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "cv2.imwrite('capture.jpg', f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "fin = np.frombuffer(bytearray(hdmi_out.frame_raw(2)), dtype=np.uint8).reshape((1080, 1920, 3))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 31,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "cv2.imwrite('capture_in.jpg', fin)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}