def __init__(self, para_p, para_r, para_o, test_img):
self.device = torch.device("cuda")
self.test_img = test_img
self.image = Image.open(test_img) # 用于抠图输入下一层
self.img = Image.open(test_img) # 复制图片用于图像金字塔
self.pnet = net.PNet().to(self.device)
self.pnet.load_state_dict(torch.load(para_p))
self.pnet.eval()
self.rnet = net.RNet().to(self.device)
self.rnet.load_state_dict(torch.load(para_r))
self.rnet.eval()
self.onet = net.ONet().to(self.device)
self.onet.load_state_dict(torch.load(para_o))
self.onet.eval()
def __init__(self, test_img):
self.device = torch.device(
"cuda" if torch.cuda.is_available() else "cpu")
self.image = Image.fromarray(np.uint8(
test_img[:, :, [2, 1, 0]])) # for croped transform cv2 to PIL
self.img = self.image # for pyramid
self.pnet = net.PNet().to(self.device)
self.pnet.load_state_dict(torch.load("./params/pnet.pkl"))
self.pnet.eval()
self.rnet = net.RNet().to(self.device)
self.rnet.load_state_dict(torch.load("./params/rnet.pkl"))
self.rnet.eval()
self.onet = net.ONet().to(self.device)
self.onet.load_state_dict(torch.load("./params/onet.pkl"))
self.onet.eval()
def __init__(self, mode: str):
"Device"
self.device = torch.device(
"cuda" if torch.cuda.is_available() else "cpu")
"Net"
self.name = mode.lower() + "net"
if mode == "P" or mode == "p":
self.net = net.PNet().to(self.device)
self.train = dataset.P_TRAIN
self.test = dataset.P_TEST
elif mode == "R" or mode == "r":
self.net = net.RNet().to(self.device)
self.train = dataset.R_TRAIN
self.test = dataset.R_TEST
elif mode == "O" or mode == "o":
self.net = net.ONet().to(self.device)
self.train = dataset.O_TRAIN
self.test = dataset.O_TEST
if os.path.exists("./params/{}.pkl".format(self.name)):
print("MODE: {} >>> Loading ... ...".format(mode.upper()))
self.net.load_state_dict(
torch.load("./params/{}.pkl".format(self.name)))
else:
print("MODE: {} >>> Initing ... ...".format(mode.upper()))
self.net.apply(functions.weights_init)
"Optimize"
self.opt = optim.Adam(self.net.parameters())
"Loss"
self.Closs = nn.BCEWithLogitsLoss()
self.Hloss = nn.BCEWithLogitsLoss(reduction='none')
self.Oloss = nn.MSELoss()
self.LMloss = nn.MSELoss()
"Tensorboard"
self.draw = SummaryWriter(
log_dir="./runs/{}_runs".format(mode.lower()))
def __init__(self, mode: str, batch_size: int):
"Device Config"
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
"Model Config"
self.mode = mode
if mode == "P" or mode == "p":
self.size = 12
self.threshold = 0.9
self.net = net.PNet().to(self.device)
elif mode == "R" or mode == "r":
self.size = 24
self.threshold = 0.99
self.net = net.RNet().to(self.device)
elif mode == "O" or mode == "o":
self.size = 48
self.threshold = 0.999
self.net = net.ONet().to(self.device)
if len(os.listdir("./params")) > 3:
print("MODE: {} >>> Loading ... ...".format(mode))
self.net.load_state_dict(torch.load("./params/{}net.pkl".format(mode.lower())))
"Dataloader Config"
self.train = data.DataLoader(dataset.choice(mode.lower()), batch_size=batch_size, shuffle=True, num_workers=4, drop_last=True)
self.test = data.DataLoader(dataset.choice("{}v".format(mode.lower())), batch_size=512, shuffle=True, drop_last=True)
"Optim Config"
optimize = optim.SGD(self.net.parameters(), lr=3e-4, momentum=0.9)
self.lookahead = Lookahead(optimize, k=5, alpha=0.5)
# self.lr = optim.lr_scheduler.CosineAnnealingLR(self.lookahead, T_max=1, eta_min=1e-5, last_epoch=-1)
"Loss Config"
self.loss_confi = nn.BCELoss()
self.loss_resolve = nn.BCELoss(reduction="none")
self.loss_offset = nn.MSELoss()
"Show Config"
self.summarywriter = SummaryWriter(log_dir="./runs/{}_runs".format(mode.lower()))
def __init__(
self,
p_net_param=r"./params_p/p_net_30.pth",
# r_net_param="./params_r/r_net_13_s.pth",
# 2(7),3(10),5(7),6(8),7(10),10(9),12(8), -13(5),15(8),21(7)
r_net_param=r"./params_r/r_net_53_s.pth",
o_net_param="./params_o/o_net_35.pth",
isCuda=False):
self.isCuda = isCuda
self.p_net = net.PNet().to(DEVICE)
self.r_net = net.RNet().to(DEVICE)
self.o_net = net.ONet().to(DEVICE)
self.p_net.load_state_dict(torch.load(p_net_param))
self.r_net.load_state_dict(torch.load(r_net_param))
self.o_net.load_state_dict(torch.load(o_net_param))
self.p_net.eval()
self.r_net.eval()
self.o_net.eval()
self.__image_transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225]) # 对于图像这个是有经验值的
])
# -*-coding:utf-8-*-
import os
import net
import torch
import trainer
model = net.PNet()
save = "F:/MTCNN/test/pnet.pth"
train = "F:/MTCNN/train"
validation = "F:/MTCNN/validation"
size = 12
if __name__ == "__main__":
trainer.Trainer(model, train, validation, save, size).main()
import net
import trainer
import os
if __name__ == '__main__':
net = net.PNet()
index = 30
last_acc = 0.9803
last_r2 = 0.8545
save_path = fr"./params_p/p_net_{index}.pth"
img_dir = r"F:\Dataset02\mtcnn_fxs"
if not os.path.exists("./params_p"):
os.makedirs("./params_p")
trainer = trainer.Trainer(net, save_path, img_dir, index, 12, last_acc,
last_r2)
trainer.train(0.7)
import net
import train
if __name__ == '__main__':
ne = net.PNet()
ne.train()
trainer = train.Trainer(ne, 'E:/save_path/12/param/', r"E:\save_path\12")
trainer.train()
