def forward(self):
self.net.train()
for step, (img, cls, off) in enumerate(self.data_loader):
img = img.cuda()
cls = cls.cuda()
off = off.cuda()
dbg(img.shape, cls.shape, off.shape, lv=dlv)
_cls, _off = self.net(img)
_cls = _cls.view(-1, 1)
_off = _off.view(-1, 4)
# Train cls
index = torch.lt(cls, 2)
# dbg(index, lv=dlv)
cls_s = cls[index]
_cls_s = _cls[index]
# dbg(cls_s, _cls_s, lv=dlv + 1)
cls_loss = self.cls_loss_fn(_cls_s, cls_s)
# Train off
index = torch.gt(cls, 0)
# dbg(index[:, 0], lv=dlv)
off_s = off[index[:, 0]]
_off_s = _off[index[:, 0]]
# dbg(off_s.shape, _off_s.shape, lv=dlv)
off_loss = self.off_loss_fn(_off_s, off_s)
loss = cls_loss + off_loss
if step % 100 == 0:
dbg("loss: {}, cls: {}, off: {}".format(
loss, cls_loss, off_loss),
lv=dlv + 1)
self.opt.zero_grad()
loss.backward()
self.opt.step()
def detect(pnet, pnet_model, img):
pnet.load_state_dict(torch.load(pnet_model))
img = np.array(img)
img = transform(img)
img = img.unsqueeze(0)
dbg("detect", img.shape, lv=dlv)
boxes = pnet_detect(pnet, img)
return boxes
def forward(self, x):
conv = self.conv(x)
dbg(conv.shape, lv=dlv)
fc = conv.view(-1, 3 * 3 * 128)
fc = self.fc(fc)
dbg(fc.shape, lv=dlv)
fc_1 = self.fc_1(fc)
fc_2 = self.fc_2(fc)
return fc_1, fc_2
def __init__(self, root, size, transform=None):
super(mtcnn_dataset, self).__init__()
dbg("init", lv=dlv)
self.transform = transform
self.labels = []
self.labels.extend(
list(open(os.path.join(root, str(size), "positive.txt"))))
self.labels.extend(
list(open(os.path.join(root, str(size), "part.txt"))))
self.labels.extend(
list(open(os.path.join(root, str(size), "negative.txt"))))
dbg(self.labels, len(self.labels), lv=dlv)
def __getitem__(self, index):
dbg("getitem", lv=dlv)
item = self.labels[index].split()
dbg(item, lv=dlv)
self.img_path = item[0]
self.cls = float(item[1])
self.off_x0 = float(item[2])
self.off_y0 = float(item[3])
self.off_x1 = float(item[4])
self.off_y1 = float(item[5])
self.cls = torch.Tensor([self.cls])
self.off = torch.Tensor(
[self.off_x0, self.off_y0, self.off_x1, self.off_y1])
dbg(self.img_path,
self.cls,
self.off_x0,
self.off_y0,
self.off_x1,
self.off_y1,
lv=dlv + 1)
img = Image.open(self.img_path)
if self.transform is not None:
dbg("Transform is not None", lv=dlv)
img = self.transform(img)
return img, self.cls, self.off
def __init__(self, net, size, model):
self.net = net.cuda() # type: nn.Module
self.model = model
self.dataset = mtcnn_dataset(save_path, size, transform)
self.data_loader = DataLoader(self.dataset,
batch_size=128,
shuffle=True,
drop_last=True)
self.cls_loss_fn = nn.MSELoss()
self.opt = optim.Adam(self.net.parameters(), lr=0.0001)
self.off_loss_fn = nn.MSELoss()
if os.path.exists(self.model):
dbg(self.model + " is exist", lv=dlv)
self.net.load_state_dict(torch.load(self.model))
def pnet_detect(net, img):
conf, off = net(img) # type: pnet
conf = conf.squeeze()
# print(torch.lt(conf, 0.5))
coor = []
box = []
height, width = conf.shape
for h in range(height):
for w in range(width):
if conf[h, w] > 0.5:
coor.append([h, w])
# print(conf[h, w].item())
print(coor)
# exit()
dbg("pd", conf.shape, off.shape, lv=dlv)
for c in coor:
# print(c[0], c[1])
box.append([
2 * c[1], 2 * c[0], 2 * c[1] + 12, 2 * c[0] + 12,
conf[c[0], c[1]].item()
])
# print(box)
return box
def __len__(self):
dbg("len", lv=dlv)
return len(self.labels)
def forward(self, x):
conv = self.conv(x)
dbg(conv.shape, lv=dlv)
conv_1 = self.conv_1(conv)
conv_2 = self.conv_2(conv)
return conv_1, conv_2
self.fc_2 = nn.Linear(256, 4)
def forward(self, x):
conv = self.conv(x)
dbg(conv.shape, lv=dlv)
fc = conv.view(-1, 3 * 3 * 128)
fc = self.fc(fc)
dbg(fc.shape, lv=dlv)
fc_1 = self.fc_1(fc)
fc_2 = self.fc_2(fc)
return fc_1, fc_2
if __name__ == '__main__':
pnet = pnet()
rnet = rnet()
onet = onet()
test = torch.randn(1, 3, 12, 12)
t1, t2 = pnet(test)
dbg(t1.shape, t2.shape, lv=dlv)
test = torch.randn(1, 3, 24, 24)
t1, t2 = rnet(test)
dbg(t1.shape, t2.shape, lv=dlv)
test = torch.randn(1, 3, 48, 48)
t1, t2 = onet(test)
dbg(t1.shape, t2.shape, lv=dlv)
img_path = "../img_celeba_4dbg"
label_path = "../img_celeba_4dbg/list_bbox_celeba.txt"
save_path = "../../img_celeba_4dbg_simple"
positive = "positive"
part = "part"
negative = "negative"
COUNT = 50
SCALE = 0.8
SIZE = [48, 24, 12]
# print(os.listdir(img_path))
for size in SIZE:
dbg("sample size:", size, lv=1)
size_save_path = os.path.join(save_path, str(size))
pos_dir = os.path.join(size_save_path, positive)
part_dir = os.path.join(size_save_path, part)
neg_dir = os.path.join(size_save_path, negative)
pos_name = pos_dir + ".txt"
part_name = part_dir + ".txt"
neg_name = neg_dir + ".txt"
for dir in [pos_dir, part_dir, neg_dir]:
if not os.path.exists(dir):
os.makedirs(dir)
dbg("dir", pos_dir, part_dir, neg_dir, lv=0)
dbg("name:", pos_name, part_name, neg_name, lv=0)