ocr-backend/config.py

from ultralytics import YOLO
from paddleocr import PaddleOCR
import platform, torch, os, json
import pydantic
os.environ['KMP_DUPLICATE_LIB_OK'] = 'TRUE'
barcode_type_list = ['RFID吊牌', '普通吊牌', '洗水唛', '饰品洗水唛', '饰品吊牌']
if platform.system().lower() == "linux":
    model_dir = "./ocr_needs/yolo-material/yolo_v8_model/"
    ocr_det = "./ocr_needs/ocr_models/en_PP-OCRv3_det_infer"
    ocr_rec = "./ocr_needs/ocr_models/en_PP-OCRv4_rec_infer"
    ocr_det_ch = "./ocr_needs/ocr_models/ch_PP-OCRv4_det_infer"
    ocr_rec_ch = "./ocr_needs/ocr_models/ch_PP-OCRv4_rec_infer"
    ocr_cls = "./ocr_needs/ocr_models/ch_ppocr_mobile_v2.0_cls_infer"
    ocr_images_dir = './ocr_workspace/'
    ttf_path = './ocr_needs/MiSans-Heavy.ttf'
header = '/data'
address = 'ocr.gloria.com.cn'
port = 8066
file_url = f'http://{address}:{port}{header}'
# 检查CUDA是否可用
device = 'cuda' if torch.cuda.is_available() else 'cpu'
print(f'device: {device} is used')

dirs_to_check = [model_dir, ocr_images_dir]
# ocr = PaddleOCR(use_angle_cls=True, lang='en', show_log=False, det_model_dir=ocr_det,
#                rec_model_dir=ocr_rec, cls_model_dir=ocr_cls)
ocr = PaddleOCR(use_angle_cls=True, lang='ch', show_log=False, det_model_dir=ocr_det_ch,
               rec_model_dir=ocr_rec_ch, cls_model_dir=ocr_cls, use_gpu=True)
print('ocr model is ready!')
for directory in dirs_to_check:
    if not os.path.exists(directory):
        os.makedirs(directory)
        print(f"Directory {directory} created.")
    else:
        print(f"Directory {directory} already exists.")

# Define the model path
model_path = os.path.join(model_dir, "obb-s-best-4.pt")
model = YOLO(model_path).to(device)
model.fuse()
print('yolo model is ready !')


sql_config = {
    'user': 'root',
    'password': 'Ywkj2020',
    'host': '10.40.0.30',
    'db': 'unique-code'
}

color_config = {
    'user': 'IT',
    'password': 'it@2022*',
    'host': '10.40.0.81',
    'database': 'goelia_erp',
}

class Search(pydantic.BaseModel):
    matio_id: str = None
    item_num : str = None
    uploadStartTime : str = None
    uploadEndTime : str = None
    ocrStartTime : str = None
    ocrEndTime : str = None
    difference : str = None
    pageSize : int = 20
    pageNum : int = 1
    barcode_type: str = None

class Matio(pydantic.BaseModel):
    barcode: str = ''
    matio_id: str = ''
    image_origin_path : str = 'None'
    upload_time : str = 'None'
    barcode_type : str = 'RFID吊牌'

class ID(pydantic.BaseModel):
    id: str = None

class Item(pydantic.BaseModel):
    name: str
    description: str