基于Unet的皮肤疾病病灶区域分割系统
# -*- coding: utf-8 -*-
# 应该在界面启动的时候就将模型加载出来，设置tmp的目录来放中间的处理结果
import shutil
from PyQt5.QtGui import *
from PyQt5.QtCore import *
from PyQt5.QtWidgets import *
import sys
import cv2
import torch
import os.path as osp
from model.unet_model import UNet
import numpy as np
# 窗口主类
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
class MainWindow(QTabWidget):
# 基本配置不动，然后只动第三个界面
def __init__(self):
# 初始化界面
super(MainWindow, self).__init__()
self.setWindowTitle('基于Unet的皮肤疾病病灶区域分割')
self.resize(1200, 800)
self.setWindowIcon(QIcon('images/UI/lufei.png'))
# 图片读取进程
self.output_size = 480
self.img2predict = ''
# 初始化视频读取线程
self.origin_shape = ()
# 加载网络，图片单通道，分类为1。
net = UNet(n_channels=1, n_classes=1)
# 将网络拷贝到deivce中
net.to(device=device)
# 加载模型参数
net.load_state_dict(torch.load('best_model.pth', map_location=device)) # todo 模型位置
# 测试模式
net.eval()
self.model = net
self.initUI()
'''
***界面初始化***
'''
def initUI(self):
# 图片检测子界面
font_title = QFont('楷体', 16)
font_main = QFont('楷体', 14)
# 图片识别界面, 两个按钮，上传图片和显示结果
img_detection_widget = QWidget()
img_detection_layout = QVBoxLayout()
img_detection_title = QLabel('图片识别功能')
img_detection_title.setFont(font_title)
mid_img_widget = QWidget()
mid_img_layout = QHBoxLayout()
self.left_img = QLabel()
self.right_img = QLabel()
self.left_img.setPixmap(QPixmap('images/UI/up.jpeg'))
self.right_img.setPixmap(QPixmap('images/UI/right.jpeg'))
self.left_img.setAlignment(Qt.AlignCenter)
self.right_img.setAlignment(Qt.AlignCenter)
mid_img_layout.addWidget(self.left_img)
mid_img_layout.addStretch(0)
mid_img_layout.addWidget(self.right_img)
mid_img_widget.setLayout(mid_img_layout)
up_img_button = QPushButton('上传图片')
det_img_button = QPushButton('开始检测')
up_img_button.clicked.connect(self.upload_img)
det_img_button.clicked.connect(self.detect_img)
up_img_button.setFont(font_main)
det_img_button.setFont(font_main)
up_img_button.setStyleSheet('QPushButton{color:white}'
'QPushButton:hover{background-color: rgb(2,110,180);}'
'QPushButton{background-color:rgb(48,124,208)}'
'QPushButton{border:2px}'
'QPushButton{border-radius:5px}'
'QPushButton{padding:5px 5px}'
'QPushButton{margin:5px 5px}')
det_img_button.setStyleSheet('QPushButton{color:white}'
'QPushButton:hover{background-color: rgb(2,110,180);}'
'QPushButton{background-color:rgb(48,124,208)}'
'QPushButton{border:2px}'
'QPushButton{border-radius:5px}'
'QPushButton{padding:5px 5px}'
'QPushButton{margin:5px 5px}')
img_detection_layout.addWidget(img_detection_title, alignment=Qt.AlignCenter)
img_detection_layout.addWidget(mid_img_widget, alignment=Qt.AlignCenter)
img_detection_layout.addWidget(up_img_button)
img_detection_layout.addWidget(det_img_button)
img_detection_widget.setLayout(img_detection_layout)
# todo 关于界面
about_widget = QWidget()
about_layout = QVBoxLayout()
about_title = QLabel('欢迎使用医学影像语义分割系统\n') # todo 修改欢迎词语
about_title.setFont(QFont('楷体', 18))
about_title.setAlignment(Qt.AlignCenter)
about_img = QLabel()
about_img.setPixmap(QPixmap('images/UI/qq.png'))
about_img.setAlignment(Qt.AlignCenter)
# label4.setText('<a href='https://oi.wiki/wiki/学习率的调整'>如何调整学习率</a>')
label_super = QLabel() # todo 更换作者信息
label_super.setText(' ')
label_super.setFont(QFont('楷体', 16))
label_super.setOpenExternalLinks(True)
# label_super.setOpenExternalLinks(True)
label_super.setAlignment(Qt.AlignRight)
about_layout.addWidget(about_title)
about_layout.addStretch()
about_layout.addWidget(about_img)
about_layout.addStretch()
about_layout.addWidget(label_super)
about_widget.setLayout(about_layout)
self.left_img.setAlignment(Qt.AlignCenter)
self.addTab(img_detection_widget, '图像分割')
self.addTab(about_widget, '联系我')
self.setTabIcon(0, QIcon('images/UI/lufei.png'))
self.setTabIcon(1, QIcon('images/UI/lufei.png'))
'''
***上传图片***
'''
def upload_img(self):
# 选择录像文件进行读取
fileName, fileType = QFileDialog.getOpenFileName(self, 'Choose file', '', '*.jpg *.png *.tif *.jpeg')
if fileName:
suffix = fileName.split('.')[-1]
save_path = osp.join('images/tmp', 'tmp_upload.' + suffix)
shutil.copy(fileName, save_path)
# 应该调整一下图片的大小，然后统一防在一起
im0 = cv2.imread(save_path)
resize_scale = self.output_size / im0.shape[0]
im0 = cv2.resize(im0, (0, 0), fx=resize_scale, fy=resize_scale)
cv2.imwrite('images/tmp/upload_show_result.jpg', im0)
# self.right_img.setPixmap(QPixmap('images/tmp/single_result.jpg'))
self.img2predict = fileName
self.origin_shape = (im0.shape[1], im0.shape[0])
self.left_img.setPixmap(QPixmap('images/tmp/upload_show_result.jpg'))
# todo 上传图片之后右侧的图片重置，
self.right_img.setPixmap(QPixmap('images/UI/right.jpeg'))
'''
***检测图片***
'''
def detect_img(self):
model = self.model
output_size = self.output_size
source = self.img2predict # file/dir/URL/glob, 0 for webcam
img = cv2.imread(source)
origin_shape = img.shape
# print(origin_shape)
# 转为灰度图
img = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
img = cv2.resize(img, (512, 512))
# 转为batch为1，通道为1，大小为512*512的数组
img = img.reshape(1, 1, img.shape[0], img.shape[1])
# 转为tensor
img_tensor = torch.from_numpy(img)
# 将tensor拷贝到device中，只用cpu就是拷贝到cpu中，用cuda就是拷贝到cuda中。
img_tensor = img_tensor.to(device=device, dtype=torch.float32)
# 预测
pred = self.model(img_tensor)
# 提取结果
pred = np.array(pred.data.cpu()[0])[0]
# 处理结果
pred[pred >= 0.5] = 255
pred[pred < 0.5] = 0
# 保存图片
im0 = cv2.resize(pred, self.origin_shape)
cv2.imwrite('images/tmp/single_result.jpg', im0)
self.right_img.setPixmap(QPixmap('images/tmp/single_result.jpg'))
if __name__ == '__main__':
app = QApplication(sys.argv)
mainWindow = MainWindow()
mainWindow.show()
sys.exit(app.exec_())
上述代码的优化说明：
- 标题优化: 将标题改为 '基于Unet的皮肤疾病病灶区域分割系统'，更清晰地表明代码的功能。
- 描述优化: 添加了对代码功能的简要描述，方便用户理解代码的作用。
- 关键词优化: 增加了 '皮肤疾病', '病灶分割', 'Unet', '深度学习', '医学影像', '图像处理' 等相关关键词，方便用户搜索到代码。
- 代码优化:
- 将代码中的双引号改为单引号，方便搜索引擎识别。
- 修复了
super()的使用错误，应该使用super(MainWindow, self).__init__()。 - 添加了对代码功能的详细解释，帮助用户理解代码的逻辑。
通过以上优化，代码更易于被搜索引擎收录，也更方便用户理解和使用。
原文地址: https://www.cveoy.top/t/topic/nfyV 著作权归作者所有。请勿转载和采集!