#Instructions:
#(script based on: https://www.pyimagesearch.com/2020/07/27/opencv-grabcut-foreground-segmentation-and-extraction/)

#run script as: python grabcut_bbox_bike.py
#when a figure appears, mouse click on the figure and press any key to continue
#this script will show a sequence of figures, and pauses until you press
#any key after a mouse click on the figure

#missing pixels are randomly selected, so run the script multiple times
#to get a sense of what outputs can be expected.

#change the percentage of missing pixels on line 41

# import the necessary packages
import numpy as np
import argparse
import time
import cv2
import os
import matplotlib.pyplot as plt

# construct the argument parser and parse the arguments
ap = argparse.ArgumentParser()
ap.add_argument("-i", "--image", type=str,
	default=os.path.sep.join(["images", "zebra2.jpg"]),
	help="path to input image that we'll apply GrabCut to")
ap.add_argument("-c", "--iter", type=int, default=40,
	help="# of GrabCut iterations (larger value => slower runtime)")
args = vars(ap.parse_args())

# load the input image from disk and then allocate memory for the
# output mask generated by GrabCut -- this mask should hae the same
# spatial dimensions as the input image
image = cv2.imread(args["image"])

#subsample
image = image[::2,::2,:]

#####################################################
#add corruption #####################################
percentage_missing = 0.10 #This number indicates the percentage of missing pixels
####################################################
####################################################

index_list = np.linspace(0,np.prod(image.shape[0:2])-1,np.prod(image.shape[0:2]),dtype=int)
np.random.shuffle(index_list)
last_ind = np.round(percentage_missing*np.prod(image.shape[0:2])).astype(int)
index_list=index_list[:last_ind]
n=image.shape
image=image.reshape(np.prod(image.shape[0:2]),3)
image[index_list,:]=0
image=image.reshape(n)

         

#plt.imshow(image[:,:,1])
print(image.shape)
mask = np.zeros(image.shape[:2], dtype="uint8")

# define the bounding box coordinates 
rect = (20, 20, 220, 300)

# allocate memory for two arrays that the GrabCut algorithm internally
# uses when segmenting the foreground from the background
fgModel = np.zeros((1, 65), dtype="float")
bgModel = np.zeros((1, 65), dtype="float")

# apply GrabCut using the the bounding box segmentation method
start = time.time()
(mask, bgModel, fgModel) = cv2.grabCut(image, mask, rect, bgModel,
	fgModel, iterCount=args["iter"], mode=cv2.GC_INIT_WITH_RECT)
# (mask, bgModel, fgModel) = cv2.grabCut(image, mask, rect, bgModel,
# 	fgModel, iterCount=20, mode=cv2.GC_INIT_WITH_RECT)
end = time.time()
print("[INFO] applying GrabCut took {:.2f} seconds".format(end - start))

# the output mask has for possible output values, marking each pixel
# in the mask as (1) definite background, (2) definite foreground,
# (3) probable background, and (4) probable foreground
values = (
	("Definite Background", cv2.GC_BGD),
	("Probable Background", cv2.GC_PR_BGD),
	("Definite Foreground", cv2.GC_FGD),
	("Probable Foreground", cv2.GC_PR_FGD),
)

# loop over the possible GrabCut mask values
for (name, value) in values:
	# construct a mask that for the current value
	print("[INFO] showing mask for '{}'".format(name))
	valueMask = (mask == value).astype("uint8") * 255
	
    # display the mask so we can visualize it
	cv2.imshow(name, valueMask)
	cv2.waitKey(0)

# we'll set all definite background and probable background pixels
# to 0 while definite foreground and probable foreground pixels are
# set to 1
outputMask = np.where((mask == cv2.GC_BGD) | (mask == cv2.GC_PR_BGD),
	0, 1)

# scale the mask from the range [0, 1] to [0, 255]
outputMask = (outputMask * 255).astype("uint8")

# apply a bitwise AND to the image using our mask generated by
# GrabCut to generate our final output image
output = cv2.bitwise_and(image, image, mask=outputMask)    

# show the input image followed by the mask and output generated by
image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

plt.figure(figsize=(5,4))
plt.imshow(outputMask);plt.title("GrabCut output ("+str(100*percentage_missing)+"% missing)");plt.savefig("GrabCutOutputMask.png")

plt.figure(figsize=(5,4))
plt.imshow(outputMask);
plt.imshow(image_rgb,alpha=0.6);plt.title("GrabCut output ("+str(100*percentage_missing)+"% missing)")
plt.axis('off')
plt.savefig("GrabCutDataPlusOutputMask.png")


# GrabCut and bitwise masking
cv2.imshow("Input", image)
cv2.waitKey(0)    
cv2.imshow("GrabCut Mask", outputMask)
cv2.imshow("GrabCut Output", output)
cv2.waitKey(0)    