#!/usr/bin/env python
"""
tests
"""
from __future__ import division, print_function
import os
import numpy as np
import cv2
import matplotlib.pyplot as plt
from PIL import Image


#img_fname = '/home/user/adapt_ws/data/DSC00084.JPG'
img_fname = '/home/local/KHQ/matt.brown/libraries/adapt/adapt_ros_ws/data/DSC00131.JPG'
mask_fname = '/home/local/KHQ/matt.brown/libraries/adapt/adapt_ros_ws/data/DSC00131_mask.png'
gif_fname = '/home/local/KHQ/matt.brown/libraries/adapt/adapt_ros_ws/data/DSC00131_auto_seg2.gif'


img = cv2.imread(img_fname)[:, :, ::-1]
mask0 = cv2.imread(mask_fname)[:, :, ::-1]

#plt.imshow(img)
#plt.imshow(mask0)

mask = np.zeros(img.shape[:2], np.uint8)
mask[:] = 2
bg_mask = np.logical_and(mask0[:, :, 0] == 255, mask0[:, :, 1] == 0)
bg_mask = np.logical_and(bg_mask, mask0[:, :, 2] == 0)
fg_mask = np.logical_and(mask0[:, :, 0] == 0, mask0[:, :, 1] == 0)
fg_mask = np.logical_and(fg_mask, mask0[:, :, 2] == 255)
mask[bg_mask] = 0
mask[fg_mask] = 1


bgdModel = np.zeros((1,65),np.float64)
fgdModel = np.zeros((1,65),np.float64)


mask, bgdModel, fgdModel = cv2.grabCut(img, mask, None, bgdModel, fgdModel, 5,
                                       cv2.GC_INIT_WITH_MASK)


im1 = Image.fromarray(img)
im2 = Image.fromarray(cv2.applyColorMap(mask*85, cv2.COLORMAP_JET))
im1.save(gif_fname, save_all=True, append_images=[im2], duration=500, loop=0)

