from tensorflow import keras
import os
import os.path
import sys
import cv2
import numpy as np

"""Module for predicting rotation using regression convolutional network
   Usage:
   > python PredictRotationCNN.py <rotationCNNfile> <inputImage> <outputFile>
   
   Input arguments:
   * <rotationCNNfile> - File storing trained rotation network.
   * <inputImage> - Input image with standardized pose of the product.
   * <outputFile> - File name for the output csv containing the predicted angle.

"""

__author__ = "Ondrej Klima"
__copyright__ = "Copyright 2020"
__credits__ = ["Ondrej Klima"]
__email__ = "iklima@fit.vutbr.cz"
__license__ = "BUT"
__version__ = "1.0"
__maintainer__ = "Ondrej Klima"


def main():
    # Parsing input arguments 
    argv = sys.argv

    try:
        cnnFileName = argv[1]
    except IndexError:
        raise IndexError('File with Mask-RCNN must be supplied as an argument')
    if not path.isfile(cnnFileName):
        raise ValueError('File "%s" does not exist!' % cnnFileName)
    
    try:
        imageFileName = argv[2]
    except IndexError:
        raise IndexError('Input image must be supplied as an argument')
    if not path.isfile(imageFileName):
        raise ValueError('File "%s" does not exist!' % imageFileName)   
    
    try:
        outputFile = argv[3]
    except IndexError:
        raise IndexError('Output direrctory for product mask must be supplied as an argument')
        

    model = keras.models.load_model(cnnFileName)
    img = cv2.imread(imageFileName)
    img = cv2.resize(img,(256,256))

    
    angle = model.predict(x=np.array(img)/255.)
    
    with open(outputFile, 'w', newline='') as csvfile:
        wr = csv.writer(csvfile, quoting=csv.QUOTE_ALL)
        wr.writerow([angle])

if __name__ == "__main__":
    main()