image_capture

#!/usr/bin/env python3

import sys
import threading
import time
import tkinter as tk
import scipy.misc
import numpy as np
import pyscreenshot

from utils import *

class RedBorder:
    def configure(self, master):
        self.window = master
        self.window.overrideredirect(True)
        self.window.resizable(width=False, height=False)
        self.window.configure(background = 'red')

    def set_pos(self, x, y, w, h):
        self.window.geometry("%dx%d+%d+%d" % (w, h, x, y))
        self.window.deiconify()

    def get_pos(self):
        return self.window.winfo_x(), self.window.winfo_y()

class CaptureWindow(RedBorder):
    def __init__(self, shape):
        self.shape = shape
        self.pos = 20, 20
        self.mouse_pos = None
        self.prev_capture = None
        self.last_moved = time.time()

    def show(self):
        self.configure(tk.Tk())
        self.window.mainloop()

    def configure(self, master):
        RedBorder.configure(self, master)
        self.borders = [self] + [self.new_border() for x in range(3)]
        self.draw_borders()
        for b in self.borders:
            self.make_draggable(b.window)

    def new_border(self):
        b = RedBorder()
        b.configure(tk.Toplevel(self.window))
        return b

    def draw_borders(self):
        self.last_moved = time.time()
        x, y = self.pos
        h, w = self.shape
        line = 10
        self.borders[0].set_pos(x-line, y-line, w+line*2, line)
        self.borders[1].set_pos(x-line, y+h, w+line*2, line)
        self.borders[2].set_pos(x-line, y, line, h)
        self.borders[3].set_pos(x+w, y, line, h)

    def make_draggable(self, w):
        w.bind("<ButtonPress-1>", self.on_button_press)
        w.bind("<ButtonRelease-1>", self.on_button_release)
        w.bind("<B1-Motion>", self.on_mouse_move)

    def on_button_press(self, event):
        self.mouse_pos = event.x, event.y
        for b in self.borders:
            b.window.config(cursor = "fleur")

    def on_button_release(self, event):
        self.on_mouse_move(event)
        self.mouse_pos = None
        for b in self.borders:
            b.window.config(cursor = "arrow")

    def on_mouse_move(self, event):
        if self.mouse_pos is None:
            return
        diff = event.x - self.mouse_pos[0], event.y - self.mouse_pos[1]
        self.pos = self.pos[0] + diff[0], self.pos[1] + diff[1]
        self.draw_borders()

def capture_image(pos, shape):
    x, y = pos
    h, w = shape[0:2]
    im = pyscreenshot.grab(bbox = (x, y, x+w, y+h))
    if shape[2:] == [3]:
        im = im.convert("RGB")
    elif shape[2:] in ([], [1]):
        im = im.convert("L")
    else:
        raise ValueError("Invalid number of channels")
    return np.asarray(im, dtype="uint8")

def capture_stream(ui, shape):
    while True:
        if time.time() < ui.last_moved + 0.3:
            time.sleep(ui.last_moved + 0.3 - time.time())
            continue

        img = capture_image(ui.pos, shape)
        if img.shape[0:2] != shape[0:2]:
            img = scipy.misc.imresize(img, shape[0:2])

        if ui.prev_capture is not None:
            diff = np.absolute(ui.prev_capture - img)
            if np.sum(diff) < 10.0:
                time.sleep(0.3)
                continue
        ui.prev_capture = img

        try:
            write_array(sys.stdout.buffer, img)
        except:
            break

def run():
    if len(sys.argv) < 2:
        sys.stderr.write("\nUsage:\n\n");
        sys.stderr.write("\timage_capture H,W\n")
        sys.stderr.write("\timage_capture H,W,3\n\n")
        sys.exit(1)

    if sys.stdout.isatty():
        sys.stderr.write("Error: Refusing to write binary data to a terminal\n")
        sys.exit(1)

    shape = str_to_image_shape(sys.argv[1])
    ui = CaptureWindow(shape[0:2])
    thread = threading.Thread(target = ui.show)
    thread.daemon = True
    thread.start()
    try:
        capture_stream(ui, shape)
    except KeyboardInterrupt:
        pass

if __name__ == "__main__":
    run()