#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Thu Feb 13 13:19:30 2020

@author: yektazahed
"""

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Wed Feb 12 14:12:38 2020

@author: yektazahed
"""

import torch
from torch.autograd import Variable
import cv2
from data import BaseTransform, VOC_CLASSES as labelmap
from ssd import build_ssd
import imageio


#create a function that gets a video and a SSD network as an input and outputs the same frame but with detection
def detect(frame, net, transform):
    
    height, width = frame.shape[:2]
    # transform the original frame to something readable for the NN
    frame_t = transform(frame)[0] #?? what is the 0?
    x = torch.from_numpy(frame_t).permute(2,0,1)
    x = Variable(x.unsqueeze(0))
    
    #give the transformed frame to the network
    y = net(x)
    #detections: [batch, category, number of occurances (score, x0, y0, x1, y1)]
    detections = y.data
    
    #create a scale that is the same size as the frame but 4D (a rectangle)
    scale = torch.Tensor([width, height, width, height])
    
    for i in range(detections.size(1)):
        
        j = 0
        
        #only plot the ones with score above 0.6
        while detections[0, i, j, 0] >= 0.6:
            pt = (detections[0, i, j, 1:]*scale).numpy()
            cv2.rectangle(frame,(int(pt[0]), int(pt[1])), (int(pt[2]), int(pt[3])), (255, 0, 0), 2)
            cv2.putText(frame, labelmap[i-1], (int(pt[0]), int(pt[1])), cv2.FONT_HERSHEY_SIMPLEX, 2, (255, 255, 255), 2, cv2.LINE_AA)
            j+=1
    
    return frame

#build the network
net = build_ssd('test')
net.load_state_dict(torch.load('ssd300_mAP_77.43_v2.pth', map_location = lambda storage, loc : storage))

#build the transform
transform = BaseTransform(net.size, (104/256.0, 117/256.0, 123/256.0))

# Detect objects from a video
#reader = imageio.get_reader('epic_horses.mp4')
#reader = imageio.get_reader('HarvardLibrary.mp4')
reader = imageio.get_reader('Hamids_Office.mp4')
fps = reader.get_meta_data()['fps']
writer = imageio.get_writer('myoutput_Hamids_Office.mp4', fps = fps)

for i, frame in enumerate(reader):
    frame = detect(frame, net.eval(), transform)
    writer.append_data(frame)
    print(i)
writer.close()


# THE END
    

#Questions:
# What is a torch tensor?
# What is a torch variable?

#Notes:
# detections: [batch, number of classes, number of occurance, (score, x0, y0, x1, y1)]
    