基于行空板的MultinomialNB古诗词分类 简单

#  -*- coding: UTF-8 -*-

# MindPlus
# Python
import sys
sys.path.append("/root/mindplus/.lib/thirdExtension/nick-pytesseract-thirdex")
# 导入所需模块
from pinpong.board import Board, Pin
import time
import cv2
import pytesseract
from PIL import Image, ImageFont, ImageDraw
import os
import numpy as np
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.naive_bayes import MultinomialNB
from sklearn.pipeline import make_pipeline
from sklearn.model_selection import train_test_split
from joblib import dump, load

# 加载预训练的机器学习模型
model = load('model.joblib')

# 初始化开发板
Board().begin()

# 定义函数在图像上绘制中文文字
def drawChinese(text, x, y, size, r, g, b, a, img):
    font = ImageFont.truetype("HYQiHei_50S.ttf", size)
    img_pil = Image.fromarray(img)
    draw = ImageDraw.Draw(img_pil)
    draw.text((x, y), text, font=font, fill=(b, g, r, a))
    frame = np.array(img_pil)
    return frame

# 设置 Tesseract OCR 的路径
pytesseract.pytesseract.tesseract_cmd = r'/usr/bin/tesseract'

# 打开摄像头并设置参数
cap = cv2.VideoCapture(0)
cap.set(cv2.CAP_PROP_FRAME_WIDTH, 240)
cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 320)
cap.set(cv2.CAP_PROP_BUFFERSIZE, 1)
cv2.namedWindow('cvwindow', cv2.WND_PROP_FULLSCREEN)
cv2.setWindowProperty('cvwindow', cv2.WND_PROP_FULLSCREEN, cv2.WINDOW_FULLSCREEN)

# 等待摄像头打开
while not cap.isOpened():
    continue
print("start!")

# 初始化识别内容变量
ShiBieNaRong = ''
img_word1 = ''
img_word2 = ''

# 主循环
while True:
    # 读取摄像头图像
    cvimg_success, img_src = cap.read()
    cvimg_h, cvimg_w, cvimg_c = img_src.shape
    cvimg_w1 = cvimg_h * 240 // 320
    cvimg_x1 = (cvimg_w - cvimg_w1) // 2
    img_src = img_src[:, cvimg_x1:cvimg_x1 + cvimg_w1]
    img_src = cv2.resize(img_src, (240, 320))
    cv2.imshow('cvwindow', img_src)

    # 检测按键输入
    key = cv2.waitKey(5)
    if key & 0xFF == ord('b'):  # 按下 "b" 键退出程序
        print("退出视频")
        break
    elif key & 0xFF == ord('a'):  # 按下 "a" 键进行文本识别和分类
        ShiBieNaRong = []
        try:
            if not os.path.exists("/root/image/pic/"):
                print("The folder does not exist, created automatically")
                os.system("mkdir -p /root/image/pic/")
        except IOError:
            print("IOError, created automatically")
            break
        # 保存图像
        cv2.imwrite("/root/image/pic/image.png", img_src)
        time.sleep(0.2)
        img = Image.open('/root/image/pic/image.png')
        # 使用 Tesseract 进行文本识别
        ShiBieNaRong = pytesseract.image_to_string(img, lang='chi_sim')
        print(ShiBieNaRong)
        # 使用机器学习模型进行分类预测
        probs = model.predict_proba([ShiBieNaRong])
        predicted_index = probs[0].argmax()
        predicted_label = model.classes_[predicted_index]
        confidence = probs[0][predicted_index]
        print(f'预测结果：{predicted_label}，置信度：{confidence}')

    # 在图像上绘制识别到的文本
    img_src = drawChinese(text=str(ShiBieNaRong), x=10, y=20, size=25, r=50, g=200, b=0, a=0, img=img_src)
    cv2.imshow('cvwindow', img_src)

# 释放摄像头并关闭窗口
cap.release()
cv2.destroyAllWindows()