import datetime
import os
import time
import ffmpeg
import torch
import cv2
import numpy as np
from multiprocessing import Process, Manager
from threading import Thread
from read_data import LoadImages, LoadStreams
import torch.backends.cudnn as cudnn
import easyocr
from tools.draw_chinese import cv2ImgAddText

class OCR():
    time_reference = datetime.datetime.now()
    counter_frame = 0
    processed_fps = 0

    def __init__(self,video_path=None):
        
        self.model = easyocr.Reader(['ch_sim','en'], gpu=True, model_storage_directory="weight/ocr/",download_enabled=False) # this needs to run only once to load the model into memory
            
        self.frame = [None] 

        if video_path is not None:
            self.video_name = video_path
        else:
            self.video_name = 'vid2.mp4'  # A default video file
            
        
        self.dataset = LoadImages(self.video_name)
        self.flag = 0
    
    def use_webcam(self, source):
        # self.dataset.release()  # Release any existing video capture
        # self.cap = cv2.VideoCapture(0)  # Open default webcam
        # print('use_webcam')
        source = source
        self.imgsz = 640
        cudnn.benchmark = True
        self.dataset = LoadStreams(source, img_size=self.imgsz)
 
    def class_to_label(self, x):
        return self.classes[int(x)]
    

           
    def get_frame(self):
 
        for im0s in self.dataset:
            # print(self.dataset.mode)
            # print(self.dataset)
            if self.dataset.mode == 'stream':
                img = im0s[0].copy()
            else:
                img = im0s.copy()

            result = self.model.readtext(img, detail = 0)
        
            
            img = cv2ImgAddText(img, 
                f'识别结果: {result}',
                10, 
                10,
                (0, 250, 0),
                20,)
            
            # cv2.putText(img, f'识别结果: {result}', (10, 50), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
            txt = f'{result}'
   
            ret, jpeg = cv2.imencode(".jpg", img)
            # print(jpeg.shape)
            

            return jpeg.tobytes(), txt

def time_synchronized():
    # pytorch-accurate time
    if torch.cuda.is_available():
        torch.cuda.synchronize()
    return time.time()