82 lines
2.3 KiB
Python
82 lines
2.3 KiB
Python
|
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()
|