← Back to branch summary

~paparazzi-uav/paparazzi/v5.0-manual

« back to all changes in this revision

Viewing changes to sw/ext/opencv_bebop/opencv/samples/python/grabcut.py

  • Committer: Paparazzi buildbot
  • Date: 2016-05-18 15:00:29 UTC
  • Revision ID: felix.ruess+docbot@gmail.com-20160518150029-e8lgzi5kvb4p7un9
Manual import commit 4b8bbb730080dac23cf816b98908dacfabe2a8ec from v5.0 branch.

Show diffs side-by-side

added added

removed removed

Lines of Context:
sw/ext/opencv_bebop/opencv/samples/python/grabcut.py
 
1
#!/usr/bin/env python
 
2
'''
 
3
===============================================================================
 
4
Interactive Image Segmentation using GrabCut algorithm.
 
5
 
 
6
This sample shows interactive image segmentation using grabcut algorithm.
 
7
 
 
8
USAGE:
 
9
    python grabcut.py <filename>
 
10
 
 
11
README FIRST:
 
12
    Two windows will show up, one for input and one for output.
 
13
 
 
14
    At first, in input window, draw a rectangle around the object using
 
15
mouse right button. Then press 'n' to segment the object (once or a few times)
 
16
For any finer touch-ups, you can press any of the keys below and draw lines on
 
17
the areas you want. Then again press 'n' for updating the output.
 
18
 
 
19
Key '0' - To select areas of sure background
 
20
Key '1' - To select areas of sure foreground
 
21
Key '2' - To select areas of probable background
 
22
Key '3' - To select areas of probable foreground
 
23
 
 
24
Key 'n' - To update the segmentation
 
25
Key 'r' - To reset the setup
 
26
Key 's' - To save the results
 
27
===============================================================================
 
28
'''
 
29
 
 
30
# Python 2/3 compatibility
 
31
from __future__ import print_function
 
32
 
 
33
import numpy as np
 
34
import cv2
 
35
import sys
 
36
 
 
37
BLUE = [255,0,0]        # rectangle color
 
38
RED = [0,0,255]         # PR BG
 
39
GREEN = [0,255,0]       # PR FG
 
40
BLACK = [0,0,0]         # sure BG
 
41
WHITE = [255,255,255]   # sure FG
 
42
 
 
43
DRAW_BG = {'color' : BLACK, 'val' : 0}
 
44
DRAW_FG = {'color' : WHITE, 'val' : 1}
 
45
DRAW_PR_FG = {'color' : GREEN, 'val' : 3}
 
46
DRAW_PR_BG = {'color' : RED, 'val' : 2}
 
47
 
 
48
# setting up flags
 
49
rect = (0,0,1,1)
 
50
drawing = False         # flag for drawing curves
 
51
rectangle = False       # flag for drawing rect
 
52
rect_over = False       # flag to check if rect drawn
 
53
rect_or_mask = 100      # flag for selecting rect or mask mode
 
54
value = DRAW_FG         # drawing initialized to FG
 
55
thickness = 3           # brush thickness
 
56
 
 
57
def onmouse(event,x,y,flags,param):
 
58
    global img,img2,drawing,value,mask,rectangle,rect,rect_or_mask,ix,iy,rect_over
 
59
 
 
60
    # Draw Rectangle
 
61
    if event == cv2.EVENT_RBUTTONDOWN:
 
62
        rectangle = True
 
63
        ix,iy = x,y
 
64
 
 
65
    elif event == cv2.EVENT_MOUSEMOVE:
 
66
        if rectangle == True:
 
67
            img = img2.copy()
 
68
            cv2.rectangle(img,(ix,iy),(x,y),BLUE,2)
 
69
            rect = (min(ix,x),min(iy,y),abs(ix-x),abs(iy-y))
 
70
            rect_or_mask = 0
 
71
 
 
72
    elif event == cv2.EVENT_RBUTTONUP:
 
73
        rectangle = False
 
74
        rect_over = True
 
75
        cv2.rectangle(img,(ix,iy),(x,y),BLUE,2)
 
76
        rect = (min(ix,x),min(iy,y),abs(ix-x),abs(iy-y))
 
77
        rect_or_mask = 0
 
78
        print(" Now press the key 'n' a few times until no further change \n")
 
79
 
 
80
    # draw touchup curves
 
81
 
 
82
    if event == cv2.EVENT_LBUTTONDOWN:
 
83
        if rect_over == False:
 
84
            print("first draw rectangle \n")
 
85
        else:
 
86
            drawing = True
 
87
            cv2.circle(img,(x,y),thickness,value['color'],-1)
 
88
            cv2.circle(mask,(x,y),thickness,value['val'],-1)
 
89
 
 
90
    elif event == cv2.EVENT_MOUSEMOVE:
 
91
        if drawing == True:
 
92
            cv2.circle(img,(x,y),thickness,value['color'],-1)
 
93
            cv2.circle(mask,(x,y),thickness,value['val'],-1)
 
94
 
 
95
    elif event == cv2.EVENT_LBUTTONUP:
 
96
        if drawing == True:
 
97
            drawing = False
 
98
            cv2.circle(img,(x,y),thickness,value['color'],-1)
 
99
            cv2.circle(mask,(x,y),thickness,value['val'],-1)
 
100
 
 
101
if __name__ == '__main__':
 
102
 
 
103
    # print documentation
 
104
    print(__doc__)
 
105
 
 
106
    # Loading images
 
107
    if len(sys.argv) == 2:
 
108
        filename = sys.argv[1] # for drawing purposes
 
109
    else:
 
110
        print("No input image given, so loading default image, ../data/lena.jpg \n")
 
111
        print("Correct Usage: python grabcut.py <filename> \n")
 
112
        filename = '../data/lena.jpg'
 
113
 
 
114
    img = cv2.imread(filename)
 
115
    img2 = img.copy()                               # a copy of original image
 
116
    mask = np.zeros(img.shape[:2],dtype = np.uint8) # mask initialized to PR_BG
 
117
    output = np.zeros(img.shape,np.uint8)           # output image to be shown
 
118
 
 
119
    # input and output windows
 
120
    cv2.namedWindow('output')
 
121
    cv2.namedWindow('input')
 
122
    cv2.setMouseCallback('input',onmouse)
 
123
    cv2.moveWindow('input',img.shape[1]+10,90)
 
124
 
 
125
    print(" Instructions: \n")
 
126
    print(" Draw a rectangle around the object using right mouse button \n")
 
127
 
 
128
    while(1):
 
129
 
 
130
        cv2.imshow('output',output)
 
131
        cv2.imshow('input',img)
 
132
        k = 0xFF & cv2.waitKey(1)
 
133
 
 
134
        # key bindings
 
135
        if k == 27:         # esc to exit
 
136
            break
 
137
        elif k == ord('0'): # BG drawing
 
138
            print(" mark background regions with left mouse button \n")
 
139
            value = DRAW_BG
 
140
        elif k == ord('1'): # FG drawing
 
141
            print(" mark foreground regions with left mouse button \n")
 
142
            value = DRAW_FG
 
143
        elif k == ord('2'): # PR_BG drawing
 
144
            value = DRAW_PR_BG
 
145
        elif k == ord('3'): # PR_FG drawing
 
146
            value = DRAW_PR_FG
 
147
        elif k == ord('s'): # save image
 
148
            bar = np.zeros((img.shape[0],5,3),np.uint8)
 
149
            res = np.hstack((img2,bar,img,bar,output))
 
150
            cv2.imwrite('grabcut_output.png',res)
 
151
            print(" Result saved as image \n")
 
152
        elif k == ord('r'): # reset everything
 
153
            print("resetting \n")
 
154
            rect = (0,0,1,1)
 
155
            drawing = False
 
156
            rectangle = False
 
157
            rect_or_mask = 100
 
158
            rect_over = False
 
159
            value = DRAW_FG
 
160
            img = img2.copy()
 
161
            mask = np.zeros(img.shape[:2],dtype = np.uint8) # mask initialized to PR_BG
 
162
            output = np.zeros(img.shape,np.uint8)           # output image to be shown
 
163
        elif k == ord('n'): # segment the image
 
164
            print(""" For finer touchups, mark foreground and background after pressing keys 0-3
 
165
            and again press 'n' \n""")
 
166
            if (rect_or_mask == 0):         # grabcut with rect
 
167
                bgdmodel = np.zeros((1,65),np.float64)
 
168
                fgdmodel = np.zeros((1,65),np.float64)
 
169
                cv2.grabCut(img2,mask,rect,bgdmodel,fgdmodel,1,cv2.GC_INIT_WITH_RECT)
 
170
                rect_or_mask = 1
 
171
            elif rect_or_mask == 1:         # grabcut with mask
 
172
                bgdmodel = np.zeros((1,65),np.float64)
 
173
                fgdmodel = np.zeros((1,65),np.float64)
 
174
                cv2.grabCut(img2,mask,rect,bgdmodel,fgdmodel,1,cv2.GC_INIT_WITH_MASK)
 
175
 
 
176
        mask2 = np.where((mask==1) + (mask==3),255,0).astype('uint8')
 
177
        output = cv2.bitwise_and(img2,img2,mask=mask2)
 
178
 
 
179
    cv2.destroyAllWindows()