# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# +
import cv2
import os
import numpy as np
import pandas as pd
from skimage import io
from skimage.transform import resize
import matplotlib.pyplot as plt


from func.postprocessing import find_cntr_condition, crf, find_cntr
from func.tool import get_fname
from func.plot import plt_result
os.environ["CUDA_VISIBLE_DEVICES"] = '6'
# -

# # (1) sup+cntr (2)+crf (3) +cntr 

FLICKR = 'data/ori/tesri/'
UnsupCV2DIR = 'model/Unsup_rmbg/result_sample/predict_img_postprocessd/'
MASKDIR = 'model/Unet_rmbg/181128103535/Predict_mask/'
SAVEDIR = 'data/processed/finaluse/'
masks = os.listdir(MASKDIR)
mask_path = [os.path.join(MASKDIR, i) for i in masks]

masks = os.listdir(MASKDIR)
flickr_path = [os.path.join(FLICKR, get_fname (i)+'.jpg') for i in masks]
new_path = [os.path.join(MASKDIR, i) for i in masks]

# + code_folding=[]
for i in range(len(new_path)):
# for i in range(1):
    # get path
    img_path = flickr_path[i]
    fname = get_fname(img_path)
    msk_path = new_path[i]
    
    # prepare data to [0,255], channel=3
    flickr_img = io.imread(img_path)
    flickr_img = resize(flickr_img,(256,256))
    flickr_img_255 = (flickr_img)*255
    flickr_img_255 =  flickr_img_255.astype(np.uint8)
    
    mask = io.imread(msk_path, as_grey=True)   
    if not mask.max()>1:
        mask = mask*255
    mask = mask.astype(np.uint8)
    mask3 = np.stack((mask,mask,mask),axis=2)
    
    # unsup2sup
    unsup2sup_mask = find_cntr_condition(mask3, condition=62000) 
    bk_unsup2sup_img =  ((unsup2sup_mask/255)*flickr_img_255).astype(np.uint8)    
    unsup2sup_img = ((unsup2sup_mask/255)*flickr_img_255 + 255-unsup2sup_mask).astype(np.uint8)
    
    # crf, return only 'R' channel
    crf_output = crf(flickr_img_255, mask3) 
    crf_mask = crf_output[:,:,0] +mask3[:,:,0]
    crf_mask = np.repeat(crf_mask,repeats=3).reshape((256,256,3))
    bk_crf_img = ((crf_mask/255)*flickr_img_255).astype(np.uint8)    
    crf_img = ((crf_mask/255)*flickr_img_255 + 255-crf_mask).astype(np.uint8)
    
    # find contour, return mask [h,w,3]
#     cntr_mask = find_cntr(crf_mask)
    cntr_mask = find_cntr_condition(crf_mask, condition=62000) 
    cntr_mask_01 = cntr_mask/255
    bk_cntr_img = flickr_img*cntr_mask_01
    cntr_img = ((cntr_mask/255)*flickr_img_255+ 255-cntr_mask).astype(np.uint8)
    
    # show unsup+cv2 result
    unsupcv2_img = io.imread(os.path.join(UnsupCV2DIR, fname+'.png'))
    
    # save out
    p_img  = [flickr_img, unsupcv2_img, unsup2sup_img, crf_img, cntr_img]
    p_title= ['Original', 'unsup', 'unsup2sup', 'crf', 'crf_cv2']
    fig = plt_result(p_img, p_title)
    
    save_to = os.path.join('data/processed/finalstep/', 'checking')
    if not os.path.exists(save_to):
        os.makedirs(save_to)
    fig.savefig(os.path.join(save_to, fname+'.png'), dpi=100, format='png',bbox_inches='tight')
    
    
    
    save_to = os.path.join('data/processed/finalstep/', 'unsup2sup_cntr')
    if not os.path.exists(save_to):
        os.makedirs(save_to)
    io.imsave(os.path.join(save_to, fname+'.png'), unsup2sup_img)
    save_to = os.path.join('data/processed/finalstep/', 'mask_unsup2sup_cntr')
    if not os.path.exists(save_to):
        os.makedirs(save_to)
    io.imsave(os.path.join(save_to, fname+'.png'), unsup2sup_mask)
    
    
    
    
    save_to = os.path.join('data/processed/finalstep/', 'unsup2sup_crf')
    if not os.path.exists(save_to):
        os.makedirs(save_to)
    io.imsave(os.path.join(save_to, fname+'.png'), crf_img)
    save_to = os.path.join('data/processed/finalstep/', 'mask_unsup2sup_crf')
    if not os.path.exists(save_to):
        os.makedirs(save_to)
    io.imsave(os.path.join(save_to, fname+'.png'), crf_mask)
    
    
    
    
    save_to = os.path.join('data/processed/finalstep/', 'unsup2sup_crf_cntr')
    if not os.path.exists(save_to):
        os.makedirs(save_to)
    io.imsave(os.path.join(save_to, fname+'.png'), cntr_img)
    save_to = os.path.join('data/processed/finalstep/', 'mask_unsup2sup_crf_cntr')
    if not os.path.exists(save_to):
        os.makedirs(save_to)
    io.imsave(os.path.join(save_to, fname+'.png'), cntr_mask)
# -



StarcoderdataJupyterScriptsDedupFiltered