import os
import shutil
import random
from tqdm import tqdm
import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms, models
from torch.utils.data import DataLoader
from sklearn.model_selection import train_test_split
test_dir = '/media/wagnchogn/data_disk/artifact/revise_cla_normal_artifact/dataset_test'
# 定义数据转换
data_transforms = transforms.Compose([
transforms.Resize((256, 256)),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])
# 定义设备
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
# 加载预训练的ResNet18模型
model = models.resnet18(pretrained=False)
num_ftrs = model.fc.in_features
model.fc = nn.Linear(num_ftrs, 1) # 二分类任务输出1个节点
model.load_state_dict(torch.load('best_model_weights.pth'))
model = model.to(device)
# 定义损失函数和优化器
criterion = nn.BCEWithLogitsLoss()
def get_pro(test_dir, model, device, data_transforms, patch_size=32):
img_predictions = []
imgs = os.listdir(test_dir)
for img in tqdm(imgs):
# 进行预测
input = data_transforms(img).to(device)
model.eval()
with torch.no_grad():
output = model(input)
pred = torch.sigmoid(output)
if pred <= 0.1:
img_predictions.append((img, pred))
return img_predictions
test111111111
最新推荐文章于 2024-09-07 22:26:58 发布