蛮有用 2022-05-29 13:27 采纳率: 75%
浏览 898
已结题

unet网络test测试图片全黑

用unet网络做图像语义分割,结果预测图全黑。

img

测试出的所有的图片为全黑色,全都只有1.1KB。这是怎么回事啊?

正文
1.用作训练的原图片大小为512*512--单通道

img

2.用作训练的label图大小为512*512--单通道

img

3.用作测试的test图大小为512*512--单通道

img

4.测试出的test图大小为512*512--单通道(1.1kb)全黑!

img

训练代码:

from model.unet_model import UNet
from utils.dataset import ISBI_Loader
from torch import optim
import torch.nn as nn
import torch
torch.cuda.current_device()

def train_net(net, device, data_path, epochs=30, batch_size=2, lr=0.000001):
    # 加载训练集
    isbi_dataset = ISBI_Loader(data_path)
    train_loader = torch.utils.data.DataLoader(dataset=isbi_dataset,
                                               batch_size=batch_size,
                                               shuffle=False)
    # 定义RMSprop算法
    optimizer = optim.RMSprop(net.parameters(), lr=lr, weight_decay=1e-8, momentum=0.9)
    # 定义Loss算法
    criterion = nn.BCEWithLogitsLoss()
    # best_loss统计,初始化为正无穷
    best_loss = float('inf')
    # 训练epochs次
    for epoch in range(epochs):
        # 训练模式
        net.train()
        # 按照batch_size开始训练
        for image, label in train_loader:
            optimizer.zero_grad()
            # 将数据拷贝到device中
            image = image.to(device=device, dtype=torch.uint8)
            label = label.to(device=device, dtype=torch.uint8)
            # 使用网络参数,输出预测结果
            pred = net(image)
            # 计算loss
            loss = criterion(pred, label)
            print('Loss/train', loss.item())
            # 保存loss值最小的网络参数
            if loss < best_loss:
                best_loss = loss
                torch.save(net.state_dict(), 'best_model.pth')
            # 更新参数
            loss.backward()
            optimizer.step()


if __name__ == "__main__":
    # 选择设备,有cuda用cuda,没有就用cpu

    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    # 加载网络,图片单通道1,分类为1。
    net = UNet(n_channels=1, n_classes=1)
    # 将网络拷贝到deivce中
    net.to(device=device)
    # 指定训练集地址,开始训练
    data_path = "/mnt/pro6/data/train/"
    train_net(net, device, data_path)


预测代码:

import glob
import numpy as np
import torch
import os
import cv2
from model.unet_model import UNet


if __name__ == "__main__":
    # 选择设备,有cuda用cuda,没有就用cpu
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    # 加载网络,图片单通道,分类为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))
    # 测试模式
    net.eval()
    # 读取所有图片路径
    tests_path = glob.glob('/Users/yym/projects/data/test/*.jpg')

    # 遍历所有图片
    for test_path in tests_path:
        # 保存结果地址
        save_res_path = test_path.split('.')[0] + '_res.png'
        # 读取图片
        img = cv2.imread(test_path)
        # 转为灰度图
        img = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
        # 转为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 = net(img_tensor)
        # 提取结果
        pred = np.array(pred.data.cpu()[0])[0]
        # 处理结果
        pred[pred >= 0.5] = 255
        pred[pred < 0.5] = 0
        # 保存图片
        cv2.imwrite(save_res_path, pred)



  • 写回答

0条回答 默认 最新

    报告相同问题?

    问题事件

    • 系统已结题 6月6日
    • 创建了问题 5月29日

    悬赏问题

    • ¥15 想要写一个跟百度网盘一样文件管理器,打开全部文件夹后 ,一级一级返回
    • ¥15 悬赏Python-playwright部署在centos7上
    • ¥15 psoc creator软件有没有人能远程安装啊
    • ¥15 快速扫描算法求解Eikonal方程咨询
    • ¥20 我的是道格手机,重置后屏幕右上角出现红色字的未写入tee key 和未写入google key请问怎么去掉啊
    • ¥30 关于R语言运行分区生存模型中的问题!
    • ¥15 校内二手商品转让网站
    • ¥20 高德地图聚合图层MarkerCluster聚合多个点,但是ClusterData只有其中部分数据,原因应该是有经纬度重合的地方点,现在我想让ClusterData显示所有点的信息,如何实现?
    • ¥100 求Web版SPC控制图程序包调式
    • ¥20 指导如何跑通以下两个Github代码