"""This file is to run the model inference here's the command 
python run_inference.py -i trainval/images/image_000000001.jpg -m model/model.pt"""

# import the necessary packages
import argparse
import cv2
import numpy as np
from PIL import Image
import torch
from torchvision import transforms
import config
from utils import get_model_instance_segmentation

# construct the argument parse and parse the arguments
ap = argparse.ArgumentParser()
ap.add_argument("-i", "--input_image", required=True,
    help="path to input image")
ap.add_argument("-m", "--model", required=True,
    help="path to trained pytorch model")
ap.add_argument("-c", "--confidence", type=float, default=0.85,
    help="minimum probability to filter weak detections")

args = vars(ap.parse_args())


model_path = args["model"]
input_image = args["input_image"]
confidence = args["confidence"]

# classes which our model will detect and the color object of the bounding box it will create
CLASSES=["Background","Person","Car"]

# reading the image with pillow and converion into the numpy arrays
img = Image.open(input_image)
open_cv_image = cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR)

# pytorch will work on the suitable device wheather it's CPU or GPU
device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")

# getting the model instance and loading the pytorch model
model = get_model_instance_segmentation(config.num_classes)
model.load_state_dict(torch.load(model_path))

# move model to the right device
model.to(device)
model.eval()

trans =transforms.Compose([transforms.ToTensor()])
img = trans(img).cuda()

# getting the all the detections generated by the trained model
detections = model([img])

# seperating out all the bounding boxes, labels, and scores we get
_bboxes, _labels, _scores = detections[0]['boxes'], detections[0]['labels'], detections[0]['scores']

# loop over the detections
for i in range(0, len(_bboxes)):
    # extract the confidence (i.e., probability) associated with the
    # prediction
    pred_confidence = _scores[i]
    # filter out weak detections by ensuring the confidence is
    # greater than the minimum confidence
    if pred_confidence > confidence:
        # extract the index of the class label from the detections,
        # then compute the (x, y)-coordinates of the bounding box
        # for the object
        idx = int(_labels[i])
        box = _bboxes[i].detach().cpu().numpy()
        (startX, startY, endX, endY) = box.astype("int")
        # display the prediction to our terminal
        label = "{}: {:.2f}%".format(CLASSES[idx], pred_confidence * 100)
        print("[INFO] {}".format(label))
        
        # draw the bounding box and label on the image      
        cv2.rectangle(open_cv_image, (startX, startY), (endX, endY),
            (0,0,255) if idx==1 else (0,255,0), 1)     
        y = startY - 15 if startY - 15 > 15 else startY + 15        
        cv2.putText(open_cv_image, label, (startX, y),
            cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0,0,255) if idx==1 else (0,255,0), 2)


# show the output image
cv2.imshow("output", open_cv_image) 
cv2.waitKey(0)
