def main(resume, use_cuda=False, use_augment=False):
## path
if True:
timestamp = datetime.now().strftime(r"%Y-%m-%d_%H-%M-%S")
save_path = os.path.join('detection', timestamp)
if not os.path.exists(save_path):
os.makedirs(save_path)
print('make a test result folder: ', save_path)
else:
save_path = None
## cuda or cpu
if use_cuda:
device = torch.device("cuda:0")
print("using cuda")
else:
device = torch.device("cpu")
print("using CPU")
if use_augment:
print("data are augmented randomly")
## dataloader
test_path = './metadata/test_images.json'
new_test_path = './metadata/new_test_images.json'
detection_path = './metadata/detection_test_images.json'
dataset = SeqDataset(phase='test',
do_augmentations=use_augment,
metafile_path=detection_path,
return_gt_label=False)
data_loader = DataLoader(
dataset,
batch_size=1,
num_workers=1,
shuffle=False,
pin_memory=True,
)
## CNN model
output_dim = 3
model = MyNet(output_dim)
## resume a ckpt
checkpoint = torch.load(resume)
model.load_state_dict(checkpoint['state_dict'])
print(model)
## evaluate
log = evaluate(model,
data_loader,
device,
draw_path=save_path,
use_conf=True)
def predict(img):
"""
加载模型和模型预测
主要步骤:
1.加载模型(请加载你认为的最佳模型)
2.图片处理
3.用加载的模型预测图片的类别
:param img: PIL.Image 对象
:return: string, 模型识别图片的类别,
共 'cardboard','glass','metal','paper','plastic','trash' 6 个类别
"""
# 加载模型,加载请注意 model_path 是相对路径, 与当前文件同级。
# 如果你的模型是在 results 文件夹下的 dnn.h5 模型，则 model_path = 'results/dnn.h5'
model_path = "weights/mynet-64-8.16-0.81.pth"
try:
# 作业提交时测试用, 请勿删除此部分
model_path = os.path.realpath(__file__).replace('main.py', model_path)
except NameError:
model_path = './' + model_path
# -------------------------- 实现模型预测部分的代码 ---------------------------
labels = {
0: 'cardboard',
1: 'glass',
2: 'metal',
3: 'paper',
4: 'plastic',
5: 'trash'
}
trans = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])
img = trans(img)
img = torch.unsqueeze(img, dim=0)
# 加载模型
net = MyNet().cpu()
net.load_state_dict(
torch.load(model_path, map_location=torch.device('cpu')))
net.eval()
with torch.no_grad():
pred = net(img)
pred = pred.numpy()
predict = labels[np.argmax(pred)]
# -------------------------------------------------------------------------
# 返回图片的类别
return predict
def main(resume, use_cuda=False, use_augment=False):
## path
if True:
timestamp = datetime.now().strftime(r"%Y-%m-%d_%H-%M-%S")
save_path = os.path.join('detection', timestamp)
if not os.path.exists(save_path):
os.makedirs(save_path)
print('make a test result folder: ', save_path)
else:
save_path = None
## cuda or cpu
if use_cuda:
device = torch.device("cuda:0")
print("using cuda")
else:
device = torch.device("cpu")
print("using CPU")
if use_augment:
print("data are augmented randomly")
## dataloader
data_loader = StreamingDataloader(imwidth=224)
## CNN model
output_dim = 3
model = MyNet(output_dim)
## resume a ckpt
checkpoint = torch.load(resume)
model.load_state_dict(checkpoint['state_dict'])
model.eval()
model = model.to(device)
print(model)
## init a detector
detector = Detector(model, data_loader, device)
## perform detection
"""
real-time feeding the image path to the detector
"""
data_folder = '/home/yanglei/codes/WSOL/detection/small_test'
image_path = os.path.join(data_folder, 'new_test0000.png')
detector(image_path, draw_path=save_path)