graduation-project/svm_test.py

124 lines
4.6 KiB
Python

import cv2
import numpy
import time
import mediapipe as mp
import math
k = 0
def binary_picture(image):
# 灰度图像
gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
# 二值图像
ret, binary = cv2.threshold(gray, 50, 255, cv2.THRESH_BINARY)
return binary
def resize_picture(image, scale):
width = int(image.shape[1] * scale)
height = int(image.shape[0] * scale)
image = cv2.resize(image, (width, height), cv2.INTER_AREA)
return image
def draw_min_rotated_rect(image, contours):
rects = []
for i in range(len(contours)):
rect = cv2.minAreaRect(contours[i])
if rect[1][0] * rect[1][1] < 500:
continue
box = cv2.boxPoints(rect)
box = numpy.int0(box)
rects.append(rect)
cv2.drawContours(image, [box], 0, (255, 0, 0), 1)
return image, rects
def draw_min_rect(image, image_binary, contours, j):
rects = []
for i in range(len(contours)):
x, y, w, h = cv2.boundingRect(contours[i])
if w * h < 500:
continue
rect = []
if x != 0 and y != 0 and w != image.shape[1] and h != image.shape[0]:
# 左上角坐标和右下角坐标
# 如果执行里面的这个画框,就是分别来画的,
# cv2.rectangle(image, (x, y), (x + w, y + h), (0, 255, 0), 1)
rect.append(x)
rect.append(y)
rect.append(x + w)
rect.append(y + h)
rects.append(rect)
rect_image = image_binary[y:y + h, x:x + w]
rect_image = cv2.resize(rect_image, (28, 28), cv2.INTER_AREA)
# cv2.imwrite(save_path + str(i) + '.jpg',hv_flip)
# cv2.imwrite('sample\\draw\\false\\' + str(j) + '.jpg', rect_image)
j = j + 1
return image, rects, j
def judge_control(rects, image_binary):
global k
for i in range(len(rects)):
k = k + 1
one_rect = rects[i]
rect_image = image_binary[one_rect[1]:one_rect[3], one_rect[0]:one_rect[2]]
rect_image = cv2.resize(rect_image, (28, 28), cv2.INTER_AREA)
save_image = rect_image
rect_image = rect_image.astype(numpy.float32)
rect_image = rect_image.reshape(-1, )
rect_image = rect_image.reshape(1, -1)
cv2.normalize(rect_image, rect_image)
page_down_svm = cv2.ml.SVM_load('PAGE_DOWN_SVM.xml')
page_up_svm = cv2.ml.SVM_load('PAGE_UP_SVM.xml')
ppt_begin_svm = cv2.ml.SVM_load('PPT_BEGIN_SVM.xml')
ppt_end_svm = cv2.ml.SVM_load('PPT_END_SVM.xml')
open_pencil_svm = cv2.ml.SVM_load('OPEN_PENCIL_SVM.xml')
draw_svm = cv2.ml.SVM_load('DRAW_SVM.xml')
page_down_image_predict = page_down_svm.predict(rect_image)
page_up_image_predict = page_up_svm.predict(rect_image)
ppt_begin_image_predict = ppt_begin_svm.predict(rect_image)
ppt_end_image_predict = ppt_end_svm.predict(rect_image)
open_pencil_image_predict = open_pencil_svm.predict(rect_image)
draw_image_predict = draw_svm.predict(rect_image)
# if page_up_image_predict[1][0]:
# cv2.imwrite('sample\\page_up\\true\\' + str(k) + '.jpg', save_image)
# return 1
# # else:
# # cv2.imwrite('sample\\page_up\\false\\' + str(k) + '.jpg', save_image)
# if page_down_image_predict[1][0]:
# cv2.imwrite('sample\\page_down\\true\\' + str(k) + '.jpg', save_image)
# return 2
# # else:
# # cv2.imwrite('sample\\page_down\\false\\' + str(k) + '.jpg', save_image)
# if ppt_begin_image_predict[1][0]:
# cv2.imwrite('sample\\ppt_begin\\true\\' + str(k) + '.jpg', save_image)
# return 3
# # else:
# # cv2.imwrite('sample\\ppt_begin\\false\\' + str(k) + '.jpg', save_image)
# if ppt_end_image_predict[1][0]:
# cv2.imwrite('sample\\ppt_end\\true\\' + str(k) + '.jpg', save_image)
# return 4
# if open_pencil_image_predict[1][0]:
# cv2.imwrite('sample\\open_pencil\\true\\' + str(k) + '.jpg', save_image)
# return 5
# if draw_image_predict[1][0]:
# cv2.imwrite('sample\\draw\\true\\' + str(k) + '.jpg', save_image)
# return 6, [one_rect[0], one_rect[1]]
if page_up_image_predict[1][0]:
return 1
if page_down_image_predict[1][0]:
return 2
if ppt_begin_image_predict[1][0]:
return 3
if ppt_end_image_predict[1][0]:
return 4
if open_pencil_image_predict[1][0]:
return 5
# if draw_image_predict[1][0]:
# return 6
return 0