#!/usr/bin/env python
from importlib import import_module
import os
from flask import Flask, render_template, Response
import cv2
import torch
from utils.torch_utils import *
from utils.general import check_img_size, non_max_suppression, scale_coords
from models.experimental import attempt_load
from models.yolo import Model
from utils.datasets import LoadStreams, LoadImages
import random
from utils.plots import plot_one_box

app = Flask(__name__)

@app.route('/')
def index():
    """Video streaming home page."""
    return render_template('index.html')


def stream():
    # WEIGHT = 'runs/train/yolov7-bsw2/weights/best.pt'
    WEIGHT = '/root/Public/model_ssai_fine_20221209/ipark_1208_1305.pt'
    # model = torch.load_state_dict(WEIGHT, map_location=select_device('0'))
    model = attempt_load(WEIGHT, map_location='cuda:0').half()
    # model = TracedModel(model, select_device('0'), img_size=640).half()

    imgsz = 640
    stride = int(model.stride.max())  # model stride
    imgsz = check_img_size(imgsz, s=stride)  # check img_size
    device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')

    # Run inference
    model(torch.zeros(1, 3, imgsz, imgsz).to(device).type_as(next(model.parameters())))  # run once
    # Get names and colors
    names = model.module.names if hasattr(model, 'module') else model.names
    colors = [[random.randint(0, 255) for _ in range(3)] for _ in names]
    old_img_w = old_img_h = imgsz
    old_img_b = 1

    """Video streaming generator function."""
    # cap = cv2.VideoCapture('rtsp://astrodom:hdci12@192.168.170.73:554/stream1')
    dataset = LoadStreams('rtsp://astrodom:hdci12@192.168.170.73:554/stream1', img_size=imgsz, stride=stride)

    while True:
        # ret, frame = cap.read()

        # if not ret:
        #     print('error')
        #     break

        for path, img, im0s, vid_cap in dataset:
            img = torch.from_numpy(img).to(device)
            img = img.half() if True else img.float()  # uint8 to fp16/32
            img /= 255.0  # 0 - 255 to 0.0 - 1.0
            if img.ndimension() == 3:
                img = img.unsqueeze(0)

            # Warmup
            if device.type != 'cpu' and (old_img_b != img.shape[0] or old_img_h != img.shape[2] or old_img_w != img.shape[3]):
                old_img_b = img.shape[0]
                old_img_h = img.shape[2]
                old_img_w = img.shape[3]
                for i in range(3):
                    model(img, augment=True)[0]

            with torch.no_grad():   # Calculating gradients would cause a GPU memory leak
                pred = model(img, augment=True)[0]
            # print(pred)
            # print('*' * 30)

            # Apply NMS
            pred = non_max_suppression(pred, 0.40, 0.45)

            # Process detections
            for i, det in enumerate(pred):  # detections per image
                if True:  # batch_size >= 1
                    p, s, im0, frame = path[i], '%g: ' % i, im0s[i].copy(), dataset.count
                # else:
                #     p, s, im0, frame = path, '', im0s, getattr(dataset, 'frame', 0)

                p = Path(p)  # to Path
                # save_path = str(save_dir / p.name)  # img.jpg
                # txt_path = str(save_dir / 'labels' / p.stem) + ('' if dataset.mode == 'image' else f'_{frame}')  # img.txt
                # gn = torch.tensor(im0.shape)[[1, 0, 1, 0]]  # normalization gain whwh
                if len(det):
                    # Rescale boxes from img_size to im0 size
                    det[:, :4] = scale_coords(img.shape[2:], det[:, :4], im0.shape).round()

                    # Print results
                    for c in det[:, -1].unique():
                        n = (det[:, -1] == c).sum()  # detections per class
                        s += f"{n} {names[int(c)]}{'s' * (n > 1)}, "  # add to string

                    # Write results
                    for *xyxy, conf, cls in reversed(det):
                    #     if save_txt:  # Write to file
                    #         xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist()  # normalized xywh
                    #         line = (cls, *xywh, conf) if opt.save_conf else (cls, *xywh)  # label format
                    #         with open(txt_path + '.txt', 'a') as f:
                    #             f.write(('%g ' * len(line)).rstrip() % line + '\n')

                        if True:  # Add bbox to image
                            label = f'{names[int(cls)]} {conf:.2f}'
                            plot_one_box(xyxy, im0, label=label, color=colors[int(cls)], line_thickness=1)



                # x = img.permute(0,3,1,2) # (B, W, H, C) --> (B, C, W, H)
                image_bytes = cv2.imencode('.jpg', im0)[1].tobytes()
                yield (b'--frame\r\n'
                    b'Content-Type: image/jpeg\r\n\r\n' + image_bytes + b'\r\n')


@app.route('/video')
def video():
    """Video streaming route. Put this in the src attribute of an img tag."""
    return Response(stream(), content_type='multipart/x-mixed-replace; boundary=frame')


if __name__ == '__main__' :
    app.run(host='0.0.0.0', threaded=True, debug=True)