from xml.etree import ElementTree

import os

from CAM_Python import cam_bbox


def main() :
    weight = 'data/models/snapshot/googlenetCAM_malware_iter_160000.caffemodel'
    model = 'data/models/prototxt/deploy_malware_googlenetCAM.prototxt'
    ImgPath = '/media/mmlab/hdd/MalwareDetector/data/RCNN_train/Images/'
    roixmlPath = 'data/RCNN_train/Annotations/'

    net_weight, net_model, out_layer, crop_size, last_conv = cam_bbox.model_set(model='googlenet', weight_path=weight, model_path=model)
    net = cam_bbox.load_model(net_model, net_weight, 0)

    img_list = os.listdir(ImgPath)
    img_list.sort()

    start = 0
    end = len(img_list)
    for i in range(start, end) :
        print "%d/%d name : %s" %(i+1, end, img_list[i])
        txt_name = img_list[i].split(".jpg")
        txt_file_path = str(roixmlPath + txt_name[0] + '.txt')
        filename = os.path.join(ImgPath, img_list[i])

        image = cam_bbox.load_image(filename)
        image_size = image.shape

        imgfile = open(txt_file_path, 'w')

        bbox_1 = cam_bbox.cam_heatmap(net, image, out_layer, crop_size, last_conv, 0)
        bbox_2 = cam_bbox.cam_heatmap(net, image, out_layer, crop_size, last_conv, 1)

        imgfile.write("Image filename : " + img_list[i] + "\n")
        imgfile.write("Image size (X x Y x C) : " + str(image_size[0]) + " x " + str(image_size[1]) + " x " + str(image_size[2]) + "\n")
        imgfile.write("Database : \"Malware Database\"\n")
        imgfile.write("Objects with ground truth : " + str(len(bbox_1) + len(bbox_2)) + " { ")
        for i in range(len(bbox_1)):
            imgfile.write("\"malware\" ")

        for i in range(len(bbox_2)):
            imgfile.write("\"white\" ")
        imgfile.write("}\n\n")

        num = 1
        for i in range(len(bbox_1)) :
            imgfile.write("Original label for object " + str(num) + " \"malware\" : \"malware\"\n")
            imgfile.write("Center point on object " + str(num) + " \"malware\" (X, Y) : (" +
                          str(bbox_1[i][0] + int((bbox_1[i][2] - bbox_1[i][0]) / 2)) + ", " +
                          str(bbox_1[i][1] + int((bbox_1[i][3] - bbox_1[i][1]) / 2)) +")\n")
            imgfile.write("Bounding box for object " + str(num) + " \"malware\" (Xmin, Ymin) - (Xmax, Ymax) : "
                          "(" + str(bbox_1[i][0]) + ", " + str(bbox_1[i][1]) + ") - (" + str(bbox_1[i][2]) + ", " + str(bbox_1[i][3]) + ")\n\n")
            num += 1

        for i in range(len(bbox_2)) :
            imgfile.write("Original label for object " + str(num) + " \"white\" : \"white\"\n")
            imgfile.write("Center point on object " + str(num) + " \"malware\" (X, Y) : (" +
                          str(bbox_2[i][0] + int((bbox_2[i][2] - bbox_2[i][0]) / 2)) + ", " +
                          str(bbox_2[i][1] + int((bbox_2[i][3] - bbox_2[i][1]) / 2)) + ")\n")
            imgfile.write("Bounding box for object " + str(num) + " \"malware\" (Xmin, Ymin) - (Xmax, Ymax) : "
                          "(" + str(bbox_2[i][0]) + ", " + str(bbox_2[i][1]) + ") - (" + str(bbox_2[i][2]) + ", " + str(bbox_2[i][3]) + ")\n\n")
            num += 1


        imgfile.close()


if __name__ == "__main__" :
    main()
