Update area_select.py

Author: lfrothin

area_select.py

@@ -2,12 +2,14 @@
 import cv2
 
 img = cv2.imread('cafe.jpg') #import example image
-mask = np.zeros_like(cv2.imread('cafe.jpg', 0)) #black 
-
+mask = np.zeros_like(img) #black
+# red, orange, yellow, green, blue - can add more/adjust
+colors = [(0, 0, 255), (28, 172, 255), (15, 196, 241), (0, 255, 0), (255, 0, 0)]
+color_i = 0 #keep track of which color to use
 shapes = [] #list to store all drawn areas
 curr_shape_index = [] #list to store the indices of the polygon currently drawing
 def draw_polygon(event, x, y, flags, params):
-    global shapes, curr_shape_index
+    global shapes, curr_shape_index, colors, color_i
     color = (0, 0, 255) #red
     thickness = 5
     if event == cv2.EVENT_LBUTTONDOWN: #for all points except last, left click
@@ -18,7 +20,8 @@ def draw_polygon(event, x, y, flags, params):
         pts = np.array(curr_shape_index, np.int32)
 
         cv2.polylines(img, [pts], isClosed, color, thickness) #draw polygon on img
-        cv2.fillPoly(mask, [pts], 255) #fill in white where areas are
+        cv2.fillPoly(mask, [pts], colors[color_i]) #fill in where areas are w/diff colors
+        color_i = (1 + color_i) % len(colors) #to prevent index from going out of bounds
         shapes.append(curr_shape_index.copy()) #save polygon
         curr_shape_index = [] #reset current indices list