import cv2
import argparse
import numpy as np
import time
import os
#import matplotlib.pyplot as plt
# handle command line arguments
ap = argparse.ArgumentParser()
ap.add_argument('-i', '--input', required=True, help = 'path to input image')
ap.add_argument('-o', '--output', required=True, help = 'path to output dir')
ap.add_argument('-c', '--config', required=True, help = 'path to yolo config file')
ap.add_argument('-w', '--weights', required=True, help = 'path to yolo pre-trained weights')
ap.add_argument('-cl', '--classes', required=True, help = 'path to text file containing class names')
args = ap.parse_args()


vs = cv2.VideoCapture(args.input)
writer = None
(W, H) = (None, None)

# read class names from text file
classes = None
with open(args.classes, 'r') as f:
    classes = [line.strip() for line in f.readlines()]

# generate different colors for different classes 
COLORS = np.random.uniform(0, 255, size=(len(classes), 3))

# read pre-trained model and config file
net = cv2.dnn.readNet(args.weights, args.config)
ln = net.getLayerNames()
lln = [ln[i - 1] for i in net.getUnconnectedOutLayers()]

int_name = 0


# loop over frames from the video file stream
while True:
    # read the next frame from the file
    (grabbed, frame) = vs.read()
    # if the frame was not grabbed, then we have reached the end
    # of the stream
    if not grabbed:
        break
    # if the frame dimensions are empty, grab them
    if W is None or H is None:
        (H, W) = frame.shape[:2]

    #blob = cv2.dnn.blobFromImage(frame, scalefactor=1/255.0, swapRB=True)
    blob = cv2.dnn.blobFromImage(frame, scalefactor=1/255.0, size=(640, 384), swapRB=True)
    # set input blob for the network
    net.setInput(blob)
    
    start = time.time()
    outs = net.forward(lln)
    end = time.time()

    # initialization
    class_ids = []
    confidences = []
    boxes = []
    conf_threshold = 0.5
    nms_threshold = 0.4

    # for each detetion from each output layer 
    # get the confidence, class id, bounding box params
    # and ignore weak detections (confidence < 0.5)

    for out in outs:
        for detection in out:
            scores = detection[5:]
            class_id = np.argmax(scores)
            confidence = scores[class_id]
            if confidence > 0.5:
                center_x = int(detection[0] * W)
                center_y = int(detection[1] * H)
                w = int(detection[2] * W)
                h = int(detection[3] * H)
                x = center_x - w / 2
                y = center_y - h / 2
                class_ids.append(class_id)
                confidences.append(float(confidence))
                boxes.append([x, y, w, h])

    print(class_ids)

    # apply non-max suppression
    indices = cv2.dnn.NMSBoxes(boxes, confidences, conf_threshold, nms_threshold)

    # go through the detections remaining
    # after nms and draw bounding box
    for i in indices:
        box = boxes[i]
        print(box)
        x = int(box[0])
        y = int(box[1])
        w = int(box[2])
        h = int(box[3])
        
        patch = frame[y:y+h, x:x+w]
        
        if patch.size > 0:
            cv2.imwrite(os.path.join(args.output, '%05d.jpeg' % (int_name)), patch)
            int_name += 1

# release the file pointers
print("[INFO] cleaning up...")
vs.release()


