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()