inpaint.py 2.29 KB
Newer Older
wester committed
1 2 3 4
#!/usr/bin/python
import urllib2
import sys
import cv2.cv as cv
wester committed
5

wester committed
6 7 8 9 10 11
class Sketcher:
    def __init__(self, windowname, dests):
        self.prev_pt = None
        self.windowname = windowname
        self.dests = dests
        cv.SetMouseCallback(self.windowname, self.on_mouse)
wester committed
12

wester committed
13 14 15 16 17 18 19 20 21 22 23 24
    def on_mouse(self, event, x, y, flags, param):
        pt = (x, y)
        if event == cv.CV_EVENT_LBUTTONUP or not (flags & cv.CV_EVENT_FLAG_LBUTTON):
            self.prev_pt = None
        elif event == cv.CV_EVENT_LBUTTONDOWN:
            self.prev_pt = pt
        elif event == cv.CV_EVENT_MOUSEMOVE and (flags & cv.CV_EVENT_FLAG_LBUTTON) :
            if self.prev_pt:
                for dst in self.dests:
                    cv.Line(dst, self.prev_pt, pt, cv.ScalarAll(255), 5, 8, 0)
            self.prev_pt = pt
            cv.ShowImage(self.windowname, img)
wester committed
25

wester committed
26 27 28 29 30 31 32 33 34
if __name__=="__main__":
    if len(sys.argv) > 1:
        img0 = cv.LoadImage( sys.argv[1], cv.CV_LOAD_IMAGE_COLOR)
    else:
        url = 'https://raw.github.com/opencv/opencv/master/samples/c/fruits.jpg'
        filedata = urllib2.urlopen(url).read()
        imagefiledata = cv.CreateMatHeader(1, len(filedata), cv.CV_8UC1)
        cv.SetData(imagefiledata, filedata, len(filedata))
        img0 = cv.DecodeImage(imagefiledata, cv.CV_LOAD_IMAGE_COLOR)
wester committed
35

wester committed
36 37 38 39 40
    print "Hot keys:"
    print "\tESC - quit the program"
    print "\tr - restore the original image"
    print "\ti or ENTER - run inpainting algorithm"
    print "\t\t(before running it, paint something on the image)"
wester committed
41

wester committed
42 43
    cv.NamedWindow("image", 1)
    cv.NamedWindow("inpainted image", 1)
wester committed
44

wester committed
45 46 47
    img = cv.CloneImage(img0)
    inpainted = cv.CloneImage(img0)
    inpaint_mask = cv.CreateImage(cv.GetSize(img), 8, 1)
wester committed
48

wester committed
49 50 51 52
    cv.Zero(inpaint_mask)
    cv.Zero(inpainted)
    cv.ShowImage("image", img)
    cv.ShowImage("inpainted image", inpainted)
wester committed
53

wester committed
54 55 56
    sk = Sketcher("image", [img, inpaint_mask])
    while True:
        c = cv.WaitKey(0) % 0x100
wester committed
57

wester committed
58 59
        if c == 27 or c == ord('q'):
            break
wester committed
60

wester committed
61 62 63 64
        if c == ord('r'):
            cv.Zero(inpaint_mask)
            cv.Copy(img0, img)
            cv.ShowImage("image", img)
wester committed
65

wester committed
66 67 68 69
        if c == ord('i') or c == ord('\n'):
            cv.Inpaint(img, inpaint_mask, inpainted, 3, cv.CV_INPAINT_TELEA)
            cv.ShowImage("inpainted image", inpainted)
    cv.DestroyAllWindows()