import jovian
jovian.commit(environment=None)

import cv2
import time
import numpy as np
import matplotlib.pyplot as plt
from PIL import Image
from keras.preprocessing import image
from keras.models import load_model
import random 

class YOLO:

    def __init__(self, config, model, labels, size=416, confidence=0.5, threshold=0.3):
        self.confidence = confidence
        self.threshold = threshold
        self.size = size
        self.labels = labels
        self.net = cv2.dnn.readNetFromDarknet(config, model)

    def inference(self, image):
        ih, iw = image.shape[:2]

        ln = self.net.getLayerNames()
        ln = [ln[i[0] - 1] for i in self.net.getUnconnectedOutLayers()]

        blob = cv2.dnn.blobFromImage(image, 1 / 255.0, (self.size, self.size), swapRB=True, crop=False)
        self.net.setInput(blob)
        start = time.time()
        layerOutputs = self.net.forward(ln)
        end = time.time()
        inference_time = end - start

        boxes = []
        confidences = []
        classIDs = []

        for output in layerOutputs:
            for detection in output:
                scores = detection[5:]
                classID = np.argmax(scores)
                confidence = scores[classID]
                if confidence > self.confidence:
                    box = detection[0:4] * np.array([iw, ih, iw, ih])
                    (centerX, centerY, width, height) = box.astype("int")
                    x = int(centerX - (width / 2))
                    y = int(centerY - (height / 2))
                    boxes.append([x, y, int(width), int(height)])
                    confidences.append(float(confidence))
                    classIDs.append(classID)

        idxs = cv2.dnn.NMSBoxes(boxes, confidences, self.confidence, self.threshold)

        results = []
        if len(idxs) > 0:
            for i in idxs.flatten():
                x, y = (boxes[i][0], boxes[i][1])
                w, h = (boxes[i][2], boxes[i][3])
                id = classIDs[i]
                confidence = confidences[i]

                results.append((id, self.labels[id], confidence, x, y, w, h))

        return iw, ih, inference_time, results

cap=cv2.VideoCapture('vid1.mp4')
img_counter = 0
fourcc = cv2.VideoWriter_fourcc(*'XVID')
out = cv2.VideoWriter('output.mp4', fourcc, 20.0, (640,  480))
yolo = YOLO("yolov3_V1_config.cfg", "yolov3_V1_config_1000.weights", ["Poacher"])
while True:
    ret, frame = cap.read()
    img_counter += 1
    print("frame: ", img_counter)
    width, height, inference_time, results = yolo.inference(frame)
    for detection in results:
        id, name, confidence, x, y, w, h = detection
        color = (0, 255, 255)
        cv2.rectangle(frame, (x, y), (x + w, y + h), color, 2)
        cv2.putText(frame, 'Poacher', (x, y - 5), cv2.FONT_HERSHEY_SIMPLEX,0.5, color, 2)
    
    out.write(frame)
    prv = cv2.resize(frame, (500,500))
    cv2.imshow("preview", prv)
    if cv2.waitKey(1)==ord('q'):
        break
        
cap.release()
out.release()
cv2.destroyAllWindows()

frame = cv2.imread('C:\\Users\\vsneh\\Udemy-notebooks\\Poacher Detection\\test_images\\pic1.jpg')
x=frame

yolo = YOLO("yolov3_V1_config.cfg", "yolov3_V1_config_1000.weights", ["Poacher"])
width, height, inference_time, results = yolo.inference(x)

# print(results)
for detection in results:
    id, name, confidence, x, y, w, h = detection
    color = (0, 255, 255)
    cv2.rectangle(frame, (x, y), (x + w, y + h), color, 2)
    cv2.putText(frame, 'Poacher', (x, y - 5), cv2.FONT_HERSHEY_SIMPLEX,0.5, color, 2)
    
    print(x,y)
frame = cv2.resize(frame, (250,250))
cv2.imshow("preview", frame)
cv2.waitKey(0)