From fb2464d3cfc91f0a099963be9ae435b345e935de Mon Sep 17 00:00:00 2001 From: prcx7ws29 Date: Sat, 7 May 2022 23:23:05 +0800 Subject: [PATCH] add main script via git --- binary(cv2).py | 177 +++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 177 insertions(+) create mode 100644 binary(cv2).py diff --git a/binary(cv2).py b/binary(cv2).py new file mode 100644 index 0000000..9472da6 --- /dev/null +++ b/binary(cv2).py @@ -0,0 +1,177 @@ +from openvino.inference_engine import IECore +import os +import torch +import numpy as np +import time +from skimage import transform, io +from PIL import Image +from PyQt5 import QtGui +from PyQt5.QtWidgets import * +from PyQt5.QtCore import * +# import openvino.inference_engine.constants +import cv2 + +def preprocess(image_dir): + + image = cv2.imread(image_dir) # PIL格式 RGB + image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) + image = cv2.resize(image, (320, 320)) + image = np.asarray(image) + image = np.swapaxes(image, 1, 2) # 转换维度 + image = np.swapaxes(image, 0, 1) + # 归一化处理 + image[:, :, 0] = (image[:, :, 0] - np.min(image[:, :, 0])) / (np.max(image[:, :, 0]) - np.min(image[:, :, 0])) + image[:, :, 0] = (image[:, :, 0] - np.mean(image[:, :, 0])) / (np.std(image[:, :, 0])) + image[:, :, 1] = (image[:, :, 1] - np.min(image[:, :, 1])) / (np.max(image[:, :, 1]) - np.min(image[:, :, 1])) + image[:, :, 1] = (image[:, :, 1] - np.mean(image[:, :, 1])) / (np.std(image[:, :, 1])) + image[:, :, 2] = (image[:, :, 2] - np.min(image[:, :, 2])) / (np.max(image[:, :, 2]) - np.min(image[:, :, 2])) + image[:, :, 2] = (image[:, :, 2] - np.mean(image[:, :, 2])) / (np.std(image[:, :, 2])) + image = np.expand_dims(image, axis=0) # 在零维度加1 + # print("image_shape has been resized to : ", image.shape) + image = torch.from_numpy(image) # 转成tensor + return image + + +def reprocess(image, image_dir, save_dir): + image = image[:, 0, :, :] + ma = np.max(image) + mi = np.min(image) + image = (image - mi) / (ma - mi) + image = image.squeeze() + # print("image_np shape is ", image_np.shape) 320*320 + # 二值化处理 + # for h in range(320): + # for w in range(320): + # if(image_np[h][w] < 30.0/255.0): + # image_np[h][w] = 0 + # else: + # image_np[h][w] = 1 + # image = Image.fromarray(image * 255).convert('RGB') # 从numpy转成PIL格式 + # image_dir is initialized in PyQt where the user open the image + # 打开原图 + image = image * 255 + pri_image = cv2.imread(image_dir) # PIL格式 RGB + h, w = pri_image.shape[0], pri_image.shape[1] + image = cv2.resize(image, (w, h), interpolation=cv2.INTER_LINEAR) + # 保存到本地 + cv2.imwrite(save_dir, image) + # print("inferenced image has been saved at ", save_dir) + return image + + +def openimage(): + imgName, imgType = QFileDialog.getOpenFileName() + jpg = QtGui.QPixmap(imgName).scaled(label.width(), label.height()) + label.setPixmap(jpg) + # pri_image_dir = imgName + image_dir = imgName + ie = IECore() + save_dir = os.path.join(os.getcwd() + os.sep + "inference_show.png") # 保存路径 + model_name = 'u2netp(canny)' + xml_dir = os.path.join(os.getcwd(), 'saved_models', model_name + '.xml') + bin_dir = os.path.join(os.getcwd(), 'saved_models', model_name + '.bin') + net = ie.read_network(xml_dir, bin_dir) + input_blob = next(iter(net.input_info)) + out_blob = next(iter(net.outputs)) + net.batch_size = 1 + # n, c, h, w = net.input_info[input_blob].input_data.shape + # print("Model input data shape = %d, %d, %d, %d" % (n, c, h, w)) + # print("image dir = ", image_dir) + image = preprocess(image_dir=image_dir) + exec_net = ie.load_network(network=net, device_name=("CUDA" if torch.cuda.is_available() else "CPU")) + # 写界面的同学记得在选择图片后加一句 + # pri_image_dir = "用户选择打开图片的路径" + start_time = time.time() + res = exec_net.infer(inputs={input_blob: image}) + res = res[out_blob] + # print("res : ", res.shape) + # res = torch.from_numpy(res) + # 后处理 + final_img = reprocess(res, image_dir=image_dir, save_dir=save_dir) + # compute the execute time + # ---------------------------------------------------------------- + init_image = Image.open(image_dir) # RGB PIL [n, c, h, w] + + init_image = np.array(init_image) + + init_image = np.squeeze(init_image) + mask = np.squeeze(final_img) + # print("mask shape : ", mask.shape) + + img_r = init_image[:, :, 0] + img_g = init_image[:, :, 1] + img_b = init_image[:, :, 2] + # print("img_r/g/b shape : ", img_b.shape) + + mask_r = mask[:, :, 0] + mask_g = mask[:, :, 1] + mask_b = mask[:, :, 2] + # print("mask_r/g/b : ", mask_r.shape) + + height = img_r.shape[0] + weight = img_r.shape[1] + for h in range(height): + for w in range(weight): + if mask_r[h][w] <= 50: + img_r[h][w] = 255 + if mask_g[h][w] <= 50: + img_g[h][w] = 255 + if mask_b[h][w] <= 50: + img_b[h][w] = 255 + + re_img = np.stack((img_r, img_g, img_b), axis=-1) + # print("re_img shape : ", re_img.shape) + final_img = Image.fromarray(re_img) + final_save_dir = os.path.join(os.getcwd() + os.sep + "final_img.png") + final_img.save(final_save_dir) + # ---------------------------------------------------------------- + + exe_time = time.time() - start_time + print("---total execute time = ", exe_time) + ccc = save_dir # 图片路径 + jpg = QtGui.QPixmap(ccc).scaled(label.width(), label.height()) + label_2.setPixmap(jpg) + ddd = final_save_dir # 图片路径 + jpg = QtGui.QPixmap(ddd).scaled(label.width(), label.height()) + label_3.setPixmap(jpg) + # 显示时间 + a = exe_time # time为float类型变量 + label_4.setText("所用时间为：" + str(a) + "s") + +# 窗口 +app = QApplication([]) + +window = QMainWindow() +window.resize(1200, 600) +window.move(300, 310) +window.setWindowTitle('前景分割') + +#三张图片加时间显示的label +label = QLabel(window) +label.move(10,80) +label.resize(320,320) + +label_2 = QLabel(window) +label_2.move(450,80) +label_2.resize(320,320) + +label_3 = QLabel(window) +label_3.move(800,80) +label_3.resize(320,320) + +label_4 = QLabel(window) +label_4.move(450, 500) +label_4.resize(300, 50) + +button = QPushButton('RUN', window) +button.move(500, 20) +button.clicked.connect(openimage) + +button_2 = QPushButton('STOP', window) +button_2.move(1000,500) +button_2.clicked.connect(QCoreApplication.quit) + +window.show() + +app.exec_() +