def __init__(self, detection_size=(160, 160)):
     self.model = Slim()
     self.model.load_state_dict(
         torch.load(open("pretrained_weights/slim_160_latest.pth", "rb"),
                    map_location="cpu"))
     self.model.eval()
     self.model.cuda()
     self.tracker = Tracker()
     self.detection_size = detection_size
class Detector:
    def __init__(self, detection_size=(160, 160)):
        self.model = Slim()
        self.model.load_state_dict(
            torch.load(open("pretrained_weights/slim_160_latest.pth", "rb"),
                       map_location="cpu"))
        self.model.eval()
        self.model.cuda()
        self.tracker = Tracker()
        self.detection_size = detection_size

    def crop_image(self, orig, bbox):
        bbox = bbox.copy()
        image = orig.copy()
        bbox_width = bbox[2] - bbox[0]
        bbox_height = bbox[3] - bbox[1]
        face_width = (1 + 2 * 0.25) * bbox_width
        face_height = (1 + 2 * 0.25) * bbox_height
        center = [(bbox[0] + bbox[2]) // 2, (bbox[1] + bbox[3]) // 2]
        bbox[0] = max(0, center[0] - face_width // 2)
        bbox[1] = max(0, center[1] - face_height // 2)
        bbox[2] = min(image.shape[1], center[0] + face_width // 2)
        bbox[3] = min(image.shape[0], center[1] + face_height // 2)
        bbox = bbox.astype(np.int)
        crop_image = image[bbox[1]:bbox[3], bbox[0]:bbox[2], :]
        h, w, _ = crop_image.shape
        crop_image = cv2.resize(crop_image, self.detection_size)
        return crop_image, ([h, w, bbox[1], bbox[0]])

    def detect(self, img, bbox):
        crop_image, detail = self.crop_image(img, bbox)
        crop_image = (crop_image - 127.0) / 127.0
        crop_image = np.array([np.transpose(crop_image, (2, 0, 1))])
        crop_image = torch.tensor(crop_image).float().cuda()
        with torch.no_grad():
            start = time.time()
            raw = self.model(crop_image)[0].cpu().numpy()
            end = time.time()
            print("PyTorch Inference Time: {:.6f}".format(end - start))
            landmark = raw[0:136].reshape((-1, 2))
        landmark[:, 0] = landmark[:, 0] * detail[1] + detail[3]
        landmark[:, 1] = landmark[:, 1] * detail[0] + detail[2]
        landmark = self.tracker.track(img, landmark)
        _, PRY_3d = get_head_pose(landmark, img)
        return landmark, PRY_3d[:, 0]
import torch
import torch.onnx
from models.slim import Slim
from collections import OrderedDict
x = torch.randn(1, 3, 96, 96)
state_dict = torch.load("./weights/slim96_epoch_93_0.1161.pth")
new_state_dict = OrderedDict()
for k, v in state_dict.items():
    name = k[7:]  # remove `module.`
    new_state_dict[name] = v
model = Slim()
model.load_state_dict(new_state_dict)
model.eval()
torch.onnx.export(model,
                  x,
                  "./weights/slim_96_latest_93.onnx",
                  input_names=["input1"],
                  output_names=['output1'])
    avg_total_loss, avg_landmark_loss, avg_loss_pose, avg_leye_loss, avg_reye_loss, avg_mouth_loss = metrics.summary()
    print(
        "Eval Avg Loss  -- Total: {:.4f} Landmark: {:.4f} Poss: {:.4f} LEye: {:.4f} REye: {:.4f} Mouth: {:.4f}".format(
            avg_total_loss, avg_landmark_loss, avg_loss_pose, avg_leye_loss, avg_reye_loss, avg_mouth_loss))
    torch.save(model.state_dict(), open("weights/slim128_epoch_{}_{:.4f}.pth".format(epoch, avg_landmark_loss), "wb"))


if __name__ == '__main__':
    checkpoint = None
    torch.backends.cudnn.benchmark = True
    train_dataset = Landmark("train.json", input_size, True)
    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=4)
    val_dataset = Landmark("val.json", input_size, False)
    val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False, num_workers=4)

    model = Slim()
    model.train()
    model.cuda()
    if checkpoint is not None:
        model.load_state_dict(torch.load(checkpoint))
        start_epoch = int(checkpoint.split("_")[-2]) + 1
    else:
        start_epoch = 0

    wing_loss_fn = WingLoss()
    mse_loss_fn = torch.nn.MSELoss()
    bce_loss_fn = torch.nn.BCEWithLogitsLoss()

    optim = torch.optim.Adam(model.parameters(), lr=lr_value_every_epoch[0], weight_decay=5e-4)
    for epoch in range(start_epoch, 150):
        for param_group in optim.param_groups:
import torch
import torch.onnx
from models.slim import Slim

x = torch.randn(1, 3, 96, 96)
model = Slim()
model.load_state_dict(
    torch.load("./weights/slim96_epoch_37_0.1151.pth", map_location="cpu"))
model.eval()
torch.onnx.export(model,
                  x,
                  "./weights/slim_96_latest.onnx",
                  input_names=["input1"],
                  output_names=['output1'])
    checkpoint = None
    if len(sys.argv) > 1:
        checkpoint = sys.argv[1]
    torch.backends.cudnn.benchmark = True
    train_dataset = Landmark("train.json", input_size, True)
    train_loader = DataLoader(train_dataset,
                              batch_size=batch_size,
                              shuffle=True,
                              num_workers=24)
    val_dataset = Landmark("val.json", input_size, False)
    val_loader = DataLoader(val_dataset,
                            batch_size=batch_size,
                            shuffle=False,
                            num_workers=24)

    model = Slim()
    model.train()
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    if torch.cuda.device_count() > 1:
        print("Let's use", torch.cuda.device_count(), "GPUs!")
        model = nn.DataParallel(model)
    model.to(device)
    #model.cuda()

    if checkpoint is not None:
        model.load_state_dict(torch.load(checkpoint))
        start_epoch = int(checkpoint.split("_")[-2]) + 1
    else:
        start_epoch = 0

    wing_loss_fn = WingLoss()
예제 #7
0
import torch
import torch.onnx
from models.slim import Slim

x = torch.randn(1, 3, 160, 160)
model = Slim()
model.load_state_dict(
    torch.load("../pretrained_weights/slim_160_latest.pth",
               map_location="cpu"))
model.eval()
torch.onnx.export(model,
                  x,
                  "../pretrained_weights/slim_160_latest.onnx",
                  input_names=["input1"],
                  output_names=['output1'])