# import the necessary packages
from tensorflow.keras.applications.mobilenet_v2 import preprocess_input
from keras.applications.xception import Xception
from tensorflow.keras.preprocessing.image import img_to_array
from tensorflow.keras.models import load_model
import numpy as np
import argparse
import imutils
import pickle
import cv2
import torchvision.models as models
from sklearn.preprocessing import LabelBinarizer
import torchvision
import torchvision.models as models
import torchvision.transforms as transforms
from torch.autograd import Variable
import pandas as pd

imagenet_classes = pd.read_csv("Documents/Imagenet Classes.csv")

# load xception model
model = Xception(weights='imagenet', include_top=True)



# load the input image from disk
image = cv2.imread("Documents/01524.png")
image = imutils.resize(image, width=299)

# run selective search on the image to generate bounding box proposal regions
ss = cv2.ximgproc.segmentation.createSelectiveSearchSegmentation()
ss.setBaseImage(image)
ss.switchToSelectiveSearchFast()
rects = ss.process()

# initialize the list of region proposals that we'll be classifying along with their associated bounding boxes
proposals = []
boxes = []

# loop over the region proposal bounding box coordinates generated by running selective search
for (x, y, w, h) in rects[:200]:
    # extract the region from the input image, ensure size matches CNN requirements
    roi = image[y:y + h, x:x + w]
    roi = cv2.resize(roi, (299,299),interpolation=cv2.INTER_CUBIC)
    
    # further preprocess by the ROI
    roi = img_to_array(roi)
    roi = preprocess_input(roi)
    
    # update our proposals and bounding boxes lists
    proposals.append(roi)
    boxes.append((x, y, x + w, y + h))



# convert the proposals and bounding boxes into NumPy arrays
proposals_np = np.array(proposals, dtype="float32")
boxes_np = np.array(boxes, dtype="int32")

boxes_list = [list(elem) for elem in boxes]
# classify each of the proposal ROIs 
proba = model.predict(proposals_np)

labels = np.argmax(proba, axis=1).tolist()

proba_max = np.zeros((len(boxes),1))

for i in range(0,len(proba_max) - 1):
    proba_max[i,0] = proba[i,labels[i]]

proba_max = proba_max.tolist()
proba_test = []
for sublist in proba_max:
    for item in sublist:
        proba_test.append(item)
   

# test =  idxs[0].tolist()
boxes_2 = boxes_np[idxs[0]].tolist()

idxs = cv2.dnn.NMSBoxes(boxes_list, proba_test, .5,.3)