# %%
import torchvision
from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
from collections import defaultdict as ddict
import json
import torch
from torchvision import datasets, transforms as T
import numpy as np
import os
import sys
import json
import cv2
import random
from model import Model
from torchvision.utils import draw_bounding_boxes
import torch as t
import matplotlib.pyplot as plt





device = 'cpu'

model_rt_path = '/home/thebears/Seafile/Designs/ML/inaturalist_models/models/'#0210701_202822.json
newest_model = os.path.join(model_rt_path, max(os.listdir(model_rt_path)).replace('.pth',''))
with open(newest_model + '.json','r') as nmj:
    model_json = json.load(nmj)

cats = model_json['categories']
cats.sort(key=lambda x: x['new_id'])
num_cat = len(cats) + 1
model_type = model_json['model_type']
model = Model(num_cat, model_type)
labels = [x['name'] for x in cats]
model.load_state_dict(
    torch.load(newest_model + '.pth', map_location=torch.device(device))
)
model.eval()
model.to(device)






#rtdir = '/home/thebears/data/hummingbird_imagenet/hummingbird'

#ff = [x for x in os.listdir(rtdir) if x.endswith('.jpg')]

# img = os.path.join(rtdir, random.choice(ff))
# image = cv2.imread(img)[:, :, ::-1].copy()
# o = T.ToTensor()(image).to(device)
# img = o[None, :, :, :]

# ou = model(img)

#oimage = t.tensor(image, dtype=t.uint8).permute([2, 0, 1])

#matplotlib.use('Qt5Agg')




#vid_path = '/home/thebears/data/hummingbird_videos/Hummingbird_01_20210601055009.mp4'
#vid_path = '/home/thebears/data/hummingbird_videos/Hummingbird_01_20210617113038.mp4'

print('model loaded')
# %%
vid_path = '/home/thebears/data/hummingbird_videos/Hummingbird_01_20210617113038.mp4'
cap = cv2.VideoCapture(vid_path)
imgs = list()

#movie = cv2.VideoWriter('/home/thebears/Seafile/Designs/ML/inaturalist_models/output.avi',cv2.VideoWriter_fourcc('M','J','P','G'), 10, (2560,1920))

frame_num = 0
# %%
while cap.isOpened():
    ret, frame = cap.read()
    if not ret:
        break

#    if frame_num % 10 == 1:
    img = cap.read()[1]
    image = img[:, :, ::-1].copy()
    o = T.ToTensor()(image).to(device)
    img = o[None, :, :, :]

    ou = model(img)

    idx = ou[0]['labels']
    label_names = [labels[x-1] for x in idx]


    scores = ou[0]['scores']
    oimage = t.tensor(255*img.squeeze(), dtype=t.uint8)
    boxes = ou[0]['boxes']

    if boxes.shape[0] > 1:
        boxes = boxes[[1],:]
        label_names = [label_names[0]]

    if boxes.shape[0] > 0:
        label_names[0] += ' {0:0.2f}'.format(scores[0])
            
        
    ox = draw_bounding_boxes(oimage, boxes,  width=5, labels = label_names,
                             font='Victor Mono SemiBold Nerd Font Complete Mono Windows Compatible',font_size=50, fill = False, colors = (255, 255, 100, 100))
    fname = '/home/thebears/Seafile/Designs/ML/inaturalist_models/frames/frame_{0:06g}.jpg'.format(frame_num)
    from PIL import Image
    im = Image.fromarray(np.uint8(ox.permute([1,2,0]).numpy()))
    im.save(fname)

#    plt.imshow(ox.permute([1, 2, 0]))
    frame_num += 1
    print(frame_num)



# %%