plot_bbox_uid_small.py

"""
Copyright 2018 Defense Innovation Unit Experimental
All rights reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.

Modifications copyright (C) 2018 <eScience Institue at University of Washington>
Licensed under CC BY-NC-ND 4.0 License [see LICENSE-CC BY-NC-ND 4.0.markdown for details] 
Written by An Yan


"""



'''
This script is modified from process_wv.py to plot bounding
boxes with uids on chipped tiles (i.e., 256 x 256 / 512 x 512) for manual inspection 

bboxes over clouds are removed automatically by naive scripts. These 
removed uids will be recorded and written to a csv/txt file

The rest of the bboxes will be rendered to png files, which will be 
handed to manual inspection

'''



from PIL import Image
import tensorflow as tf
import io
import glob
from tqdm import tqdm
import numpy as np
import logging
import argparse
import os
import json
import wv_util as wv
import tfr_util as tfr
import aug_util as aug
import csv

"""
  A script that processes xView imagery. 
  Args:
      image_folder: A folder path to the directory storing xView .tif files
        ie ("xView_data/")

      json_filepath: A file path to the GEOJSON ground truth file
        ie ("xView_gt.geojson")

      test_percent (-t): The percentage of input images to use for test set

      suffix (-s): The suffix for output TFRecord files.  Default suffix 't1' will output
        xview_train_t1.record and xview_test_t1.record

      augment (-a): A boolean value of whether or not to use augmentation

  Outputs:
    Writes two files to the current directory containing training and test data in
        TFRecord format ('xview_train_SUFFIX.record' and 'xview_test_SUFFIX.record')
"""




def detect_blackblock(img):
    # check the # of pixels that with RGB values are all equal to 0
    w,h,c = img.shape
    black_pixel_count=0
    threshold = 0.9 * w * h * 3
    non_black_count = np.count_nonzero(img)
    if non_black_count > threshold:
        return False
    else:
        return True



def detect_clouds(img,  boxes, classes, uids):
    mean_threshold_min = 160
    w, h, _ = img.shape
    #print('w,h', w, h)

    # TODO: try the threshold, process_wv.py used 18 
    var_threshold = 25
    rows_to_delete = list()
    deleted_uids = list()
    for i in range(boxes.shape[0]):
        xmin, ymin, xmax, ymax = boxes[i]
#         ymin = 0
        if xmin < 0:
            xmin = 0
        if ymin<0:
            y_min = 0
        if xmax > h:
            print('xmax > h')
            xmax = h
        if ymax > w:
            print('ymax > w')
            ymax= h
        #print(xmin, ymin, xmax, ymax)
        # clip bbox areas
        #cropped_img = img.crop((xmin, ymin, xmax, ymax))
        cropped_img = img[int(ymin):int(ymax), int(xmin):int(xmax)]  # note the order of w/h
#         print(cropped_img)
#         print(cropped_img.shape)
        array_img =  np.array(cropped_img)
        mean_img = np.mean(array_img)
        
        #print('mean_img', mean_img, i)
        var_img = np.std(array_img)
        #print('var_img',var_img, i)
        #if var_img < var_threshold and (cropped_img> 150).all() and (cropped_img< 255).all():
        if var_img < var_threshold and mean_img > mean_threshold_min:
            print('bounding box i has cloud', i)
            # need to delete this bbox
            rows_to_delete.append(i)
            deleted_uids.append(uids[i])
    print('rows_to_delete',rows_to_delete)
            
    if len(rows_to_delete) == 0:
        return img,  boxes, classes, uids, []
    else:
        # return boxes and classes with clouds removed
        new_coords = np.delete(boxes, rows_to_delete, axis=0)
        new_classes = np.delete(classes, rows_to_delete, axis=0)
        new_uids = np.delete(uids, rows_to_delete, axis=0)
        return img, new_coords, new_classes, new_uids, deleted_uids
        
        
    



def get_images_from_filename_array(coords,chips,classes,folder_names,res=(250,250)):
    """
    Gathers and chips all images within a given folder at a given resolution.

    Args:
        coords: an array of bounding box coordinates
        chips: an array of filenames that each coord/class belongs to.
        classes: an array of classes for each bounding box
        folder_names: a list of folder names containing images
        res: an (X,Y) tuple where (X,Y) are (width,height) of each chip respectively

    Output:
        images, boxes, classes arrays containing chipped images, bounding boxes, and classes, respectively.
    """

    images =[]
    boxes = []
    clses = []

    k = 0
    bi = 0   
    
    for folder in folder_names:
        fnames = glob.glob(folder + "*.tif")
        fnames.sort()
        for fname in tqdm(fnames):
            #Needs to be "X.tif" ie ("5.tif")
            name = fname.split("\\")[-1]
            arr = wv.get_image(fname)
            
            img,box,cls = wv.chip_image(arr,coords[chips==name],classes[chips==name],res)

            for im in img:
                images.append(im)
            for b in box:
                boxes.append(b)
            for c in cls:
                clses.append(cls)
            k = k + 1
            
    return images, boxes, clses

def shuffle_images_and_boxes_classes(im,box,cls):
    """
    Shuffles images, boxes, and classes, while keeping relative matching indices

    Args:
        im: an array of images
        box: an array of bounding box coordinates ([xmin,ymin,xmax,ymax])
        cls: an array of classes

    Output:
        Shuffle image, boxes, and classes arrays, respectively
    """
    assert len(im) == len(box)
    assert len(box) == len(cls)
    
    perm = np.random.permutation(len(im))
    out_b = {}
    out_c = {}
    
    k = 0 
    for ind in perm:
        out_b[k] = box[ind]
        out_c[k] = cls[ind]
        k = k + 1
    return im[perm], out_b, out_c

'''
Datasets
_multires: multiple resolutions. Currently [(500,500),(400,400),(300,300),(200,200)]
_aug: Augmented dataset
'''

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("image_folder", help="Path to folder containing image chips (ie 'Image_Chips/' ")
    parser.add_argument("json_filepath", help="Filepath to GEOJSON coordinate file")
    parser.add_argument("-t", "--test_percent", type=float, default=0.333,
                    help="Percent to split into test (ie .25 = test set is 25% total)")
    parser.add_argument("-s", "--suffix", type=str, default='t1',
                    help="Output TFRecord suffix. Default suffix 't1' will output 'xview_train_t1.record' and 'xview_test_t1.record'")
    parser.add_argument("-a","--augment", type=bool, default=False,
    				help="A boolean value whether or not to use augmentation")
    # debug: added percent of data to produce, the purpose is to produce small dataset for fast algorithm development
    parser.add_argument("-p", "--sample_percent", type=int, default = 1, help = "Portion to sample data (1/sample_percent) from the original dataset. Meaning that only use a portion of the dataset to construct training and testing. The purpose is for fast algorithm development")
    


    args = parser.parse_args()

    logging.basicConfig(level=logging.INFO)
    logger = logging.getLogger(__name__)

    #resolutions should be largest -> smallest.  We take the number of chips in the largest resolution and make
    #sure all future resolutions have less than 1.5times that number of images to prevent chip size imbalance.
    #res = [(500,500),(400,400),(300,300),(200,200)]
    #res = [(300,300)]
    #res = [(512,512)]
    
    res = [(512, 512)]


    AUGMENT = args.augment
    # debug
    #SAVE_IMAGES = False
    SAVE_IMAGES = True
    images = {}
    boxes = {}
    train_chips = 0
    test_chips = 0


    # debug
    deleted_uids_bycloud = set()

    #Parameters
    max_chips_per_res = 100000
    train_writer = tf.python_io.TFRecordWriter("harvey_delete1_%s.record" % args.suffix)
    test_writer = tf.python_io.TFRecordWriter("harvey_delete2_%s.record" % args.suffix)

    #coords,chips,classes = wv.get_labels(args.json_filepath)
    coords,chips,classes,uids = wv.get_labels_w_uid_nondamaged(args.json_filepath)


    # debug
    #print('number of chips from geojson', len(chips))
    #print('number of classes from geojson', len(classes))
    #print('some coords: ', coords[2])
    #print('some coords: ', coords[3000])
    #print('classes some: ', classes[4])
    #print('chips  ', chips[349])

    # debug
    sample_percent = args.sample_percent
    # a list of classes to be augment. Set to set to be empty if no augmentation
    # is wanted 
    class_to_aug = set([])
    num_aug_per_class = {}  # class_id: # of augmentation generated
    for class_id in class_to_aug:
        num_aug_per_class[class_id] = 0


    #debug
    # for cloud removing and black portion removing
    num_cloud_rm = 0  # number of 512 x 512 chips that have clouds removed
    num_black = 0  # number of 512 x 512 chips that have black parts 



    for res_ind, it in enumerate(res):
        tot_box = 0
        logging.info("Res: %s" % str(it))
        ind_chips = 0

        fnames = glob.glob(args.image_folder + "*.tif")
        fnames.sort()

        for fname in tqdm(fnames):
            #Needs to be "X.tif", ie ("5.tif")
            #Be careful!! Depending on OS you may need to change from '/' to '\\'.  Use '/' for UNIX and '\\' for windows
            name = fname.split("/")[-1]
            # debug
            #print('file name: ', name)
            arr = wv.get_image(fname)
            
            # debug
            print('file name: ', name)
            #print('classes[chips==name], ', classes[chips==name])


            im,box,classes_final, uids_chip = wv.chip_image_with_uid(arr,coords[chips==name],classes[chips==name],uids[chips==name] ,it)

            #Shuffle images & boxes all at once. Comment out the line below if you don't want to shuffle images
            #im,box,classes_final = shuffle_images_and_boxes_classes(im,box,classes_final)
            split_ind = int(im.shape[0] * args.test_percent)
            
            for idx, image in enumerate(im):
                if idx%sample_percent !=0:
                    continue
                 # debug
                print('processing idx: ', idx)
                # debug
                # remove black block
                if detect_blackblock(image):
                    num_black +=1
                    continue
                # remove clouds
                
                image, new_coords, new_classes, new_uids, deleted_uids = detect_clouds(image,box[idx],classes_final[idx], uids_chip[idx])                
                if len(new_coords)!= len(box[idx]):
                    num_cloud_rm += 1
                deleted_uids_bycloud = deleted_uids_bycloud.union(set(deleted_uids))
               
                # debug: changed image,box[idx],classes_final[idx] to newly constructed img and box
               # tf_example = tfr.to_tf_example(image,box[idx],classes_final[idx])
                tf_example = tfr.to_tf_example(image, new_coords, new_classes)
                



                #Check to make sure that the TF_Example has valid bounding boxes.  
                #If there are no valid bounding boxes, then don't save the image to the TFRecord.
                float_list_value_xmin = tf_example.features.feature['image/object/bbox/xmin'].float_list.value
                float_list_value_ymin = tf_example.features.feature['image/object/bbox/ymin'].float_list.value
                float_list_value_xmax = tf_example.features.feature['image/object/bbox/xmax'].float_list.value
                float_list_value_ymax = tf_example.features.feature['image/object/bbox/ymax'].float_list.value

                if (ind_chips < max_chips_per_res and np.array(float_list_value_xmin).any() and np.array(float_list_value_xmax).any() and np.array(float_list_value_ymin).any() and np.array(float_list_value_ymax).any()):
                    tot_box+=np.array(float_list_value_xmin).shape[0]
                    
                    if idx < split_ind:
                        test_writer.write(tf_example.SerializeToString())
                        test_chips+=1
                        #if SAVE_IMAGES:
                                    # debug: changed save dir
                         #   aug.draw_bboxes(image, new_coords).save('./harvey_img_inspect_uid_val/img_%s_%s.png'%(name,str(idx)))


                    else:
                        train_writer.write(tf_example.SerializeToString())
                        train_chips += 1
                        #if SAVE_IMAGES:
                                    # debug: changed save dir
                         #   aug.draw_bboxes(image, new_coords).save('./harvey_img_inspect_uid_train/img_%s_%s.png'%(name,str(idx)))



     
                    ind_chips +=1
                    
                    # debug
                    # store the training and validation images with bboxes for inspection
                    
                    #if SAVE_IMAGES:
                                    # debug: changed save dir
                     #   aug.draw_bboxes_withindex_multiclass(image, new_coords, new_classes, new_uids).save('../harvey_manual_inspect_test_512/img_%s_%s.png'%(name,str(idx)))
                    
                


                    #Make augmentation probability proportional to chip size.  Lower chip size = less chance.
                    #This makes the chip-size imbalance less severe.
                    prob = np.random.randint(0,np.max(res))
                    #for 200x200: p(augment) = 200/500 ; for 300x300: p(augment) = 300/500 ...



                    # debug
                    # added customized data augmentation for minor classes
                    #class_to_aug = [2, 3, 4]  # damaged roads, trash heaps, and bridges
                    # Minor classes will be augmented to 63 times larger with various augmentations
                    # 1. Detect whether minor classes are in the small chips, if yes, augment 
                    # this chip. The output will be a tensor of augmented images, bboxes, and classes
                    # unpack the output to tfrecord TRAINING data. 
                    # 2. If the chip does not contain any minor classes, go to normal augmentation
                    #skip_augmentation = set()  # contains a list of chips that contain minor classes
                    MINOR_CLASS_FLAG = False
                    for class_id in class_to_aug:
                        #num_aug_per_class[class_id] = 0
                        #num_aug_this_class = 0
                        # debug
                        # print('checking whether this chip contain class: ', class_id)
                        # this chip contains minor classes
                        #if np.any(classes_final[idx][:]== class_id):
                        #if class_id in set(classes_final[idx]) and idx > split_ind:
                        if class_id in set(new_classes) and idx > split_ind:
                          #       skip_augmentation.add(idx)
                            MINOR_CLASS_FLAG = True
                         #   print('trying to call expand_aug for chip: ', idx)
                            #im_aug,boxes_aug,classes_aug= aug.expand_aug_random(image, box[idx], classes_final[idx], class_id)  
                            im_aug,boxes_aug,classes_aug= aug.expand_aug_random(image, new_coords, new_classes, class_id)  


                            #debug
                            print('augmentig chip: ', idx)
                            num_aug = 0
                            for aug_idx, aug_image in enumerate(im_aug):
                                tf_example_aug = tfr.to_tf_example(aug_image, boxes_aug[aug_idx],classes_aug[aug_idx])            
                                #Check to make sure that the TF_Example has valid bounding boxes.  
                #If there are no valid bounding boxes, then don't save the image to the TFRecord.
                                float_list_value_xmin = tf_example_aug.features.feature['image/object/bbox/xmin'].float_list.value
                                float_list_value_xmax = tf_example_aug.features.feature['image/object/bbox/xmax'].float_list.value
                                float_list_value_ymin = tf_example_aug.features.feature['image/object/bbox/ymin'].float_list.value
                                float_list_value_ymax = tf_example_aug.features.feature['image/object/bbox/ymax'].float_list.value

                                # debug
                                #num_aug = 0
                                if (np.array(float_list_value_xmin).any() and np.array(float_list_value_xmax).any() and np.array(float_list_value_ymin).any() and np.array(float_list_value_ymax).any()):
                                    tot_box+=np.array(float_list_value_xmin).shape[0]
                    
                                    train_writer.write(tf_example_aug.SerializeToString())
                                    num_aug = num_aug + 1
                                    train_chips+=1
                                    num_aug_per_class[class_id] = num_aug_per_class[class_id]+1
                         #           num_aug_this_class=num_aug_this_class + 1
                                    # debug
                                    if aug_idx%10 == 0 and SAVE_IMAGES:
                                    # debug: changed save dir
                                        aug_image = (aug_image).astype(np.uint8)
                                        aug.draw_bboxes(aug_image,boxes_aug[aug_idx]).save('./expand_aug_random_256/img_aug_%s_%s_%s_%s.png'%(name, str(idx), str(aug_idx), str(class_id)))
                            # debug
                            print('augmenting class: ', class_id)
                            print('number of augmentation: ',num_aug)
                        #num_aug_per_class[class_id] = num_aug_this_class

                    # it: iterator for different resolutions
                    # start to augment the rest
                    if AUGMENT and prob < it[0] and MINOR_CLASS_FLAG == False:
                        
                        for extra in range(3):
                            center = np.array([int(image.shape[0]/2),int(image.shape[1]/2)])
                            deg = np.random.randint(-10,10)
                            #deg = np.random.normal()*30
                            # changed
                            # remove and gaussian blur
                            
                            newimg = aug.gaussian_blur(image)
                            #newimg = image

                            #.3 probability for each of shifting vs rotating vs shift(rotate(image))
                            p = np.random.randint(0,3)
                            # debug
                            # modified to use the removed cloud version of bboxes
                            # image, new_coords, new_classes
                            if p == 0:
                                newimg,nb = aug.shift_image(newimg,new_coords)
                                #newimg,nb = aug.shift_image(newimg,box[idx])
                            elif p == 1:
                                newimg,nb = aug.rotate_image_and_boxes(newimg,deg,center,new_coords)
                                #newimg,nb = aug.rotate_image_and_boxes(newimg,deg,center,box[idx])
                            elif p == 2:
                                newimg,nb = aug.rotate_image_and_boxes(newimg,deg,center,new_coords)
                                #newimg,nb = aug.rotate_image_and_boxes(newimg,deg,center,box[idx])
                                newimg,nb = aug.shift_image(newimg,nb)
                                

                            newimg = (newimg).astype(np.uint8)

                            if idx%100 == 0 and SAVE_IMAGES:
                                #debug
                                # changed save dir
                                Image.fromarray(newimg).save('./augmented_img_60/img_%s_%s_%s.png'%(name,extra,it[0]))

                            if len(nb) > 0:
                                # debug
                                # modified to use the cloud removed bboxs
                                tf_example = tfr.to_tf_example(newimg,nb,new_classes)
                                #tf_example = tfr.to_tf_example(newimg,nb,classes_final[idx])

                                #DonI't count augmented chips for chip indices
                                # changed
                                # removed data augmentation for test data
                                if idx < split_ind:
                                  #  test_writer.write(tf_example.SerializeToString())
                                   # test_chips += 1
                                    continue
                                else:
                                    train_writer.write(tf_example.SerializeToString())
                                    train_chips+=1
                            # debug:
                	    # save image + bounding boxes for debug
                            #else:
                                    if idx%100 ==0 and SAVE_IMAGES:
                                        # debug: changed save dir
                                        aug.draw_bboxes(newimg,nb).save('./harvey_augmented/img_aug_%s_%s_%s.png'%(name,extra,it[0]))
        if res_ind == 0:
            max_chips_per_res = int(ind_chips * 1.5)
            logging.info("Max chips per resolution: %s " % max_chips_per_res)

        logging.info("Tot Box: %d" % tot_box)
        logging.info("Chips: %d" % ind_chips)

    # debug
    for key, val in num_aug_per_class.items():
        print('for class:' , key)
        print('augmentation applied: ', val)
    # debug
    print('num of black small chips removed: ', num_black)
    print('num of small chips containing clouds:', num_cloud_rm)

    logging.info("saved: %d train chips" % train_chips)
    logging.info("saved: %d test chips" % test_chips)
    train_writer.close()
    test_writer.close() 
    print('number of bboxes over clouds deleted: ', len(deleted_uids_bycloud))
    # write deleted_uids_bycloud  to file
    cloudfname = 'automatic_deleted_bbox_onclouds_' + args.suffix + '.txt'
    with open(cloudfname, 'w') as f:
        for uid in deleted_uids_bycloud:
            f.write('%s\n'% int(uid))
    f.close()