# Using digits to classify images under a specific folder

# Copyright (c) 2017, zhangys@zjgsu.edu.cn

# requires NVidia Digits to be installed

SOURCE_FOLDER = "/home/zys/seagate/Extra_20170723/"
SOURCE_FILE_LIST = "/home/zys/seagate/Extra_20170723.txt"
REF_FOLDER = '' #"/home/zys/seagate/ROP_Classified_20170603_Combined/images"+"/"  # the files with the same name in this folder will not be processed
TARGET_FOLDER = '/home/zys/seagate/Extra_Selected_20170723/' #"/home/zys/seagate/ROP_201707/A"+"/"


JOBS_DIR = "/home/zys/DIGITS-master/digits/jobs"  # "/var/lib/digits/jobs/"


import os
import sys
import shutil

caffe_root = os.environ['CAFFE_ROOT']
sys.path.insert(0,caffe_root+'python')

# Add path for DIGITS package
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))
import digits.config
from digits.inference.errors import InferenceError
from digits.job import Job
from digits import utils

# To solve error "Check failed: error == cudaSuccess (10 vs. 0)  invalid device ordinal"
import caffe
caffe.set_device(0)

"""
Perform inference on a list of images using the specified model
"""
def nv_digits_infer(input_list,
          output_dir,
          jobs_dir,
          model_id,
          epoch,
          batch_size,
          layers,
          gpu):
    """
    Perform inference on a list of images using the specified model
    """
    # job directory defaults to that defined in DIGITS config
    if jobs_dir == 'none':
        jobs_dir = digits.config.config_value('jobs_dir')

    # load model job
    model_dir = os.path.join(jobs_dir, model_id)
    assert os.path.isdir(model_dir), "Model dir %s does not exist" % model_dir
    model = Job.load(model_dir)

    # load dataset job
    dataset_dir = os.path.join(jobs_dir, model.dataset_id)
    assert os.path.isdir(dataset_dir), "Dataset dir %s does not exist" % dataset_dir
    dataset = Job.load(dataset_dir)
    for task in model.tasks:
        task.dataset = dataset

    # retrieve snapshot file
    task = model.train_task()
    snapshot_filename = None
    epoch = float(epoch)
    if epoch == -1 and len(task.snapshots):
        # use last epoch
        epoch = task.snapshots[-1][1]
        snapshot_filename = task.snapshots[-1][0]
    else:
        for f, e in task.snapshots:
            if e == epoch:
                snapshot_filename = f
                break
    if not snapshot_filename:
        raise InferenceError("Unable to find snapshot for epoch=%s" % repr(epoch))

    # retrieve image dimensions and resize mode
    image_dims = dataset.get_feature_dims()
    height = image_dims[0]
    width = image_dims[1]
    channels = image_dims[2]
    resize_mode = dataset.resize_mode if hasattr(dataset, 'resize_mode') else 'squash'

    n_input_samples = 0  # number of samples we were able to load
    input_ids = []       # indices of samples within file list
    input_data = []      # sample data

    
    # load paths from file
    paths = None
    with open(input_list) as infile:
        paths = infile.readlines()
    # load and resize images
    for idx, path in enumerate(paths):
        path = path.strip()
        try:
            image = utils.image.load_image(path.strip())            
            image = utils.image.resize_image(
                image,
                height,
                width,
                channels=channels,
                resize_mode=resize_mode)
           
            input_ids.append(idx)
            input_data.append(image)
            n_input_samples = n_input_samples + 1
        except utils.errors.LoadImageError as e:
            print e

    # perform inference

    if layers != 'none':
        raise InferenceError("Layer visualization is not supported for multiple inference")
    outputs = model.train_task().infer_many(
        input_data,
        snapshot_epoch=epoch,
        gpu=gpu,
        resize=True)
    
    return outputs["softmax"]

def nv_digits_classify(filelist, model):
    
    DIGITS_JOB_ID = DICT[model]
    softmax = nv_digits_infer(filelist, "/home/zys/data/C3R/test/tmp/",
    		  JOBS_DIR,
    		  DIGITS_JOB_ID,
    		  -1,
    		  1,
    		  'none',
    		  0)
        
    cls = []
    for idx, p in enumerate(softmax):   
        cls.append(p.argmax())
        
    i = 0
    with open(filelist) as f:
        for line in f:
            line = line.strip()            
            c = str(cls[i])
            i+=1
            directory = TARGET_FOLDER+"/" + c;
            if not os.path.exists(directory):
                os.makedirs(directory)
            shutil.copyfile(line, directory + "/" + os.path.basename(line))


def split_filelist(filelist, MAX_LINE_PER_FILE):
     i=0     
     fidx=0
     files=[]
     currentfile = None
     with open(filelist) as f:
        for line in f:
            if(i%MAX_LINE_PER_FILE == 0):
                if (currentfile is not None):
                    currentfile.close()
                currentfile = open(filelist+str(fidx), 'w')
                files.append(filelist+str(fidx))
                fidx+=1
            line = line.strip()            
            currentfile.write("%s\n" % line)
            i+=1            
     
     return files

# nv_digits_classify(SOURCE_FILE_LIST, "VGG-16")
files = split_filelist(SOURCE_FILE_LIST, 500)
i=0
import time
caffe.set_mode_gpu()

# call in the shell: Every 3 rounds, will throw error. So use for loop to handle it.
# for((i=30;i<=132;i=i+3));do sudo python ...py $i;done
print ("------",sys.argv[1],"------")
for f in files:        
    if(i >= int( sys.argv[1])):
        nv_digits_classify(f, "VGG-16")
    print(i, "*************", f, " Finished *************")
    i+=1;
    #if(i%5==0):       # In GPU mode, CUDA handle is not released immediately
    #    caffe.set_mode_gpu()
    #else:
    #    caffe.set_mode_cpu()
    #i+=1    
    time.sleep(0)