Exemple #1
0
def main(_argv):
    logger = tf.get_logger()
    logger.disabled = True
    logger.setLevel(logging.FATAL)
    set_memory_growth()

    cfg = load_yaml(FLAGS.cfg_path)

    # define network
    model = RetinaFaceModel(cfg, training=False, iou_th=FLAGS.iou_th,
                            score_th=FLAGS.score_th)

    # load checkpoint
    checkpoint_dir = FLAGS.weights
    checkpoint = tf.train.Checkpoint(model=model)
    if tf.train.latest_checkpoint(checkpoint_dir):
        checkpoint.restore(tf.train.latest_checkpoint(checkpoint_dir))
        print("[*] load ckpt from {}.".format(
            tf.train.latest_checkpoint(checkpoint_dir)))
    else:
        print("[*] Cannot find ckpt from {}.".format(checkpoint_dir))
        exit()

    model.summary()
    for i in model.layers:
        print(i.output)
    model.save(FLAGS.output)
Exemple #2
0
def initialize():
    # init
    os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
    os.environ['CUDA_VISIBLE_DEVICES'] = FLAGS.gpu

    # Setup logger
    logger = tf.get_logger()
    logger.disabled = True
    logger.setLevel(logging.FATAL)
    set_memory_growth()

    cfg = load_yaml(FLAGS.cfg_path)

    # Define network
    model = RetinaFaceModel(cfg, training=False,
                            iou_th=FLAGS.iou_th, score_th=FLAGS.score_th)

    # Load checkpoints
    checkpoint_dir = './checkpoints/' + cfg['sub_name']
    checkpoint = tf.train.Checkpoint(model=model)

    if tf.train.latest_checkpoint(checkpoint_dir):
        checkpoint.restore(tf.train.latest_checkpoint(checkpoint_dir))
        print(
            "[*] load ckpt from {}.".format(tf.train.latest_checkpoint(checkpoint_dir)))
    else:
        print("[*] Cannot find ckpt from {}.".format(checkpoint_dir))
        exit()

    return model, cfg
Exemple #3
0
def main(args):
    ijbc_meta = np.load(args.meta_path)

    os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
    os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu

    logger = tf.get_logger()
    logger.disabled = True
    logger.setLevel(logging.FATAL)
    set_memory_growth()

    #cfg = load_yaml('configs/arc_res50.yaml')
    cfg = load_yaml(args.config_path)

    model = ArcFaceModel(size=cfg['input_size'],
                         backbone_type=cfg['backbone_type'],
                         training=False)

    ckpt_path = tf.train.latest_checkpoint('./checkpoints/' + cfg['sub_name'])
    if ckpt_path is not None:
        print("[*] load ckpt from {}".format(ckpt_path))
        model.load_weights(ckpt_path)
    else:
        print("[*] Cannot find ckpt from {}.".format(ckpt_path))
        exit()

    img_names = [
        os.path.join(args.input_path,
                     img_name.split('/')[-1])
        for img_name in ijbc_meta['img_names']
    ]

    embedding_size = cfg['embd_shape']
    batch_size = cfg['batch_size']
    img_size = cfg['input_size']

    def read_img(filename):
        raw = tf.io.read_file(filename)
        img = tf.image.decode_jpeg(raw, channels=3)
        img = tf.cast(img, tf.float32)
        img = img / 255
        return img

    dataset = tf.data.Dataset.from_tensor_slices(img_names)
    dataset = dataset.map(read_img,
                          num_parallel_calls=tf.data.experimental.AUTOTUNE)
    dataset = dataset.batch(batch_size)
    dataset = dataset.prefetch(buffer_size=tf.data.experimental.AUTOTUNE)
    embeddings = model.predict(dataset, batch_size=batch_size, verbose=1)

    print('embeddings', embeddings.shape)
    np.save(args.output_path, embeddings)
Exemple #4
0
def main(_argv):
    # init
    os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
    os.environ['CUDA_VISIBLE_DEVICES'] = FLAGS.gpu

    logger = tf.get_logger()
    logger.disabled = True
    logger.setLevel(logging.FATAL)
    set_memory_growth()

    cfg = load_yaml(FLAGS.cfg_path)

    # load dataset
    test_dataset = load_cifar10_dataset(
        cfg['val_batch_size'], split='test', shuffle=False,
        drop_remainder=False, using_crop=False, using_flip=False,
        using_cutout=False)


    # define network
    # TODO : change cfg
    for num_arch in range(50):
        model = CifarModel(cfg, training=False)
        model.summary(line_length=80)
        print("param size = {:f}MB".format(count_parameters_in_MB(model)))

        # load checkpoint
        checkpoint_path = './checkpoints/' + cfg['sub_name'] + '/best.ckpt'
        try:
            model.load_weights('./checkpoints/' + cfg['sub_name'] + '/best.ckpt')
            print("[*] load ckpt from {}.".format(checkpoint_path))
        except:
            print("[*] Cannot find ckpt from {}.".format(checkpoint_path))
            exit()

        # inference
        top1 = AvgrageMeter()
        top5 = AvgrageMeter()
        for step, (inputs, labels) in enumerate(test_dataset):
            # run model
            logits = model(inputs)

            # cacludate top1, top5 acc
            prec1, prec5 = accuracy(logits.numpy(), labels.numpy(), topk=(1, 5))
            n = inputs.shape[0]
            top1.update(prec1, n)
            top5.update(prec5, n)

            print(" {:03d}: top1 {:f}, top5 {:f}".format(step, top1.avg, top5.avg))

        print("Test Acc: top1 {:.2f}%, top5 {:.2f}%".format(top1.avg, top5.avg))
def main(_):
    os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
    os.environ['CUDA_VISIBLE_DEVICES'] = FLAGS.gpu
    logger = tf.get_logger()
    logger.disabled = True
    logger.setLevel(logging.FATAL)
    set_memory_growth()

    cfg = load_yaml(FLAGS.cfg_path)
	cfg['num_classes'] = 179721
	cfg['num_samples'] = 8621403
	cfg['train_dataset'] = 'data/popet/dtat/ms1m_asian.tfrecord'
	cfg['sub_name'] = '200324_model'
	print(cfg)
Exemple #6
0
def main(_argv):
    os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
    os.environ['CUDA_VISIBLE_DEVICES'] = FLAGS.gpu

    logger = tf.get_logger()
    logger.disabled = True
    logger.setLevel(logging.FATAL)
    set_memory_growth()

    cfg = load_yaml(FLAGS.cfg_path)

    model = ArcFaceModel(size=cfg['input_size'],
                         backbone_type=cfg['backbone_type'],
                         training=False)

    ckpt_path = tf.train.latest_checkpoint('../checkpoints/' + cfg['sub_name'])
    if ckpt_path is not None:
        print("[*] load ckpt from {}".format(ckpt_path))
        model.load_weights(ckpt_path)
    else:
        print("[*] Cannot find ckpt from {}.".format(ckpt_path))
        exit()

    if FLAGS.img_path:
        str1 = 'this'
        str2 = 'this1'
        #print(str1 in str2)

        file_dir = "newTWICE_id"
        output = glob.glob(file_dir + "/*.npy")

        for i in range(100):
            sampleList = random.sample(output, 2)

            embedding1 = np.load(sampleList[0])
            embedding2 = np.load(sampleList[1])
            dist = distance(embedding1, embedding2, 1)

            str1 = sampleList[0].split("\\")[1].split(".")[0]
            str2 = sampleList[1].split("\\")[1].split(".")[0]
            print(str1)
            if dist < 0.4:
                print(str1 + " and " + str2 + "is same")
                print(dist)

                #             tp = tp + 1
            else:
                print(str1 + " and " + str2 + "is diff")
                print(dist)
Exemple #7
0
def get_vect_face_img(align_face_img):  # 벡터화된 이미지 값을 직접 넣도록 했습니다.
    os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
    os.environ['CUDA_VISIBLE_DEVICES'] = '0'  # 기본값 0

    logger = tf.get_logger()
    logger.disabled = True
    logger.setLevel(logging.FATAL)
    set_memory_growth()

    img = align_face_img
    #     img = cv2.imread(img_path) # opencv 로 이미지 읽어옴
    # img = cv2.resize(img, (cfg['input_size'], cfg['input_size'])) # 이미지 크기 조정
    img = img.astype(np.float32) / 255.  # 0. ~ 255. 사이 값 변환
    if len(img.shape) == 3:  # 차원 조절
        img = np.expand_dims(img, 0)
    embeds = l2_norm(model(img))  # l2 normalization?
    return embeds
Exemple #8
0
def main(_argv):
    os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
    os.environ['CUDA_VISIBLE_DEVICES'] = FLAGS.gpu

    logger = tf.get_logger()
    logger.disabled = True
    logger.setLevel(logging.FATAL)
    set_memory_growth()

    cfg = load_yaml(FLAGS.cfg_path)

    model = ArcFaceModel(size=cfg['input_size'],
                         backbone_type=cfg['backbone_type'],
                         training=False)

    ckpt_path = tf.train.latest_checkpoint('./checkpoints/' + cfg['sub_name'])
    if ckpt_path is not None:
        print("[*] load ckpt from {}".format(ckpt_path))
        model.load_weights(ckpt_path)
    else:
        print("[*] Cannot find ckpt from {}.".format(ckpt_path))
        exit()

    print("[*] Loading LFW, AgeDB30 and CFP-FP...")
    lfw, agedb_30, cfp_fp, lfw_issame, agedb_30_issame, cfp_fp_issame = \
        get_val_data(cfg['test_dataset'])

    print("[*] Perform Evaluation on LFW...")
    acc_lfw, best_th = perform_val(
        cfg['embd_shape'], cfg['batch_size'], model, lfw, lfw_issame,
        is_ccrop=cfg['is_ccrop'])
    print("    acc {:.4f}, th: {:.2f}".format(acc_lfw, best_th))

    print("[*] Perform Evaluation on AgeDB30...")
    acc_agedb30, best_th = perform_val(
        cfg['embd_shape'], cfg['batch_size'], model, agedb_30,
        agedb_30_issame, is_ccrop=cfg['is_ccrop'])
    print("    acc {:.4f}, th: {:.2f}".format(acc_agedb30, best_th))

    print("[*] Perform Evaluation on CFP-FP...")
    acc_cfp_fp, best_th = perform_val(
        cfg['embd_shape'], cfg['batch_size'], model, cfp_fp, cfp_fp_issame,
        is_ccrop=cfg['is_ccrop'])
    print("    acc {:.4f}, th: {:.2f}".format(acc_cfp_fp, best_th))
def main(_argv):
    os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
    os.environ['CUDA_VISIBLE_DEVICES'] = FLAGS.gpu

    logger = tf.get_logger()
    logger.disabled = True
    logger.setLevel(logging.FATAL)
    set_memory_growth()

    cfg = load_yaml(FLAGS.cfg_path)

    model = tf.keras.models.load_model(cfg['model_path'])

    print("[*] Loading LFW, AgeDB30 and CFP-FP...")
    lfw, agedb_30, cfp_fp, lfw_issame, agedb_30_issame, cfp_fp_issame = \
        get_val_data(cfg['test_dataset'])

    print("[*] Perform Evaluation on LFW...")
    acc_lfw, best_th = perform_val(cfg['embd_shape'],
                                   cfg['batch_size'],
                                   model,
                                   lfw,
                                   lfw_issame,
                                   is_ccrop=cfg['is_ccrop'])
    print("    acc {:.4f}, th: {:.2f}".format(acc_lfw, best_th))

    print("[*] Perform Evaluation on AgeDB30...")
    acc_agedb30, best_th = perform_val(cfg['embd_shape'],
                                       cfg['batch_size'],
                                       model,
                                       agedb_30,
                                       agedb_30_issame,
                                       is_ccrop=cfg['is_ccrop'])
    print("    acc {:.4f}, th: {:.2f}".format(acc_agedb30, best_th))

    print("[*] Perform Evaluation on CFP-FP...")
    acc_cfp_fp, best_th = perform_val(cfg['embd_shape'],
                                      cfg['batch_size'],
                                      model,
                                      cfp_fp,
                                      cfp_fp_issame,
                                      is_ccrop=cfg['is_ccrop'])
    print("    acc {:.4f}, th: {:.2f}".format(acc_cfp_fp, best_th))
    def __create_model(self, cfg_path):
        os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
        os.environ['CUDA_VISIBLE_DEVICES'] = '0'

        logger = tf.get_logger()
        logger.disabled = True
        # logger.setLevel(logging.FATAL)
        set_memory_growth()
        cfg = load_yaml(cfg_path)
        model = RetinaFaceModel(cfg, training=False, iou_th=0.4, score_th=0.5)
        # load checkpoint
        checkpoint_dir = './checkpoints/' + cfg['sub_name']
        checkpoint = tf.train.Checkpoint(model=model)
        if tf.train.latest_checkpoint(checkpoint_dir):
            checkpoint.restore(tf.train.latest_checkpoint(checkpoint_dir))
            print("[*] load ckpt from {}.".format(
                tf.train.latest_checkpoint(checkpoint_dir)))
        else:
            print("[*] Cannot find ckpt from {}.".format(checkpoint_dir))
            exit()
        return model
Exemple #11
0
def main(_argv):
    # init
    os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
    os.environ['CUDA_VISIBLE_DEVICES'] = FLAGS.gpu

    logger = tf.get_logger()
    logger.disabled = True
    logger.setLevel(logging.FATAL)
    set_memory_growth()

    # define network
    print("[*] Loading weights from {FLAGS.model}")
    model = tf.keras.models.load_model(FLAGS.model)
    # evaluation
    if FLAGS.down_up:
        print("[*] Processing on single image {}".format(FLAGS.img_path))
        raw_img = cv2.imread(FLAGS.img_path)
        lr_img, hr_img = create_lr_hr_pair(raw_img, FLAGS.scale)

        sr_img = tensor2img(model(lr_img[np.newaxis, :] / 255))
        bic_img = imresize_np(lr_img, FLAGS.scale).astype(np.uint8)

        str_format = "[{}] PSNR/SSIM: SR={:.2f}db/{:.2f}"
        print(
            str_format.format(
                os.path.basename(FLAGS.img_path),
                calculate_psnr(rgb2ycbcr(sr_img), rgb2ycbcr(hr_img)),
                calculate_ssim(rgb2ycbcr(sr_img), rgb2ycbcr(hr_img))))
        result_img_path = './Bic_SR_HR_' + os.path.basename(FLAGS.img_path)
        print("[*] write the result image {}".format(result_img_path))
        results_img = np.concatenate((bic_img, sr_img, hr_img), 1)
        cv2.imwrite(result_img_path, results_img)
    else:
        print("[*] Processing on single image {}".format(FLAGS.img_path))
        lr_img = cv2.imread(FLAGS.img_path)
        sr_img = tensor2img(model(lr_img[np.newaxis, :] / 255))

        result_img_path = './SR_' + os.path.basename(FLAGS.img_path)
        print("[*] write the result image {}".format(result_img_path))
        cv2.imwrite(result_img_path, sr_img)
import cv2
import numpy as np
import tensorflow as tf
import tensorflow_datasets as tfds
from modules.utils import set_memory_growth
from modules.dataset import load_cifar10_dataset

set_memory_growth()

dataset = load_cifar10_dataset(batch_size=1,
                               split='train',
                               shuffle=False,
                               using_normalize=False)

for (img, labels) in dataset:
    img = img.numpy()[0]
    print(img.shape, labels.shape, labels.numpy())

    cv2.imshow('img', cv2.cvtColor(img, cv2.COLOR_RGB2BGR))
    if cv2.waitKey(0) == ord('q'):
        exit()
Exemple #13
0
def main(_argv):
    # init
    os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
    os.environ['CUDA_VISIBLE_DEVICES'] = FLAGS.gpu

    logger = tf.get_logger()
    logger.disabled = True
    logger.setLevel(logging.FATAL)
    set_memory_growth()

    cfg = load_yaml(FLAGS.cfg_path)

    # define network
    model = RetinaFaceModel(cfg,
                            training=False,
                            iou_th=FLAGS.iou_th,
                            score_th=FLAGS.score_th)

    # load model from weights.h5
    # model.load_weights('./model/mbv2_weights.h5', by_name=True, skip_mismatch=True)

    # load checkpoint
    checkpoint_dir = './checkpoints/' + cfg['sub_name']
    checkpoint = tf.train.Checkpoint(model=model)
    if tf.train.latest_checkpoint(checkpoint_dir):
        checkpoint.restore(tf.train.latest_checkpoint(checkpoint_dir))
        print("[*] load ckpt from {}.".format(
            tf.train.latest_checkpoint(checkpoint_dir)))
    else:
        print("[*] Cannot find ckpt from {}.".format(checkpoint_dir))
        exit()

    if not FLAGS.webcam:
        file_path = '/Users/lichaochao/Downloads/images_UMU/'
        for file_name in os.listdir(file_path + 'source_images/'):
            image_path = file_path + 'source_images/' + file_name
            if not os.path.exists(image_path):
                print(f"cannot find image path from {image_path}")
                continue

            img_raw = cv2.imread(image_path)
            img = np.float32(img_raw.copy())

            # img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

            # pad input image to avoid unmatched shape problem
            img, pad_params = pad_input_image(img, max_steps=max(cfg['steps']))
            img_height, img_width, _ch = img.shape

            # run model
            outputs = model(img[np.newaxis, ...])

            preds = tf.concat([
                outputs[0][0], outputs[1][0, :, 1][..., tf.newaxis],
                outputs[2][0, :, 1][..., tf.newaxis]
            ], -1)

            priors = prior_box_tf((img_height, img_width), cfg['min_sizes'],
                                  cfg['steps'], cfg['clip'])
            decode_preds = decode_tf(preds, priors, cfg['variances'])

            selected_indices = tf.image.non_max_suppression(
                boxes=decode_preds[:, :4],
                scores=decode_preds[:, -1],
                max_output_size=tf.shape(decode_preds)[0],
                iou_threshold=FLAGS.iou_th,
                score_threshold=FLAGS.score_th)

            outputs = tf.gather(decode_preds, selected_indices).numpy()

            # recover padding effect
            outputs = recover_pad_output(outputs, pad_params)
            has_face = False
            is_smile = False
            for prior_index in range(len(outputs)):
                ann = outputs[prior_index]
                if ann[-1] >= 0.5:
                    has_face = True
                    x1, y1 = int(ann[0] * img_width), int(ann[1] * img_height)
                    x2, y2 = int(ann[2] * img_width), int(ann[3] * img_height)

                    text = "face: {:.2f}".format(ann[-1] * 100)
                    cv2.putText(img, text, (x1 + 5, y1 - 10),
                                cv2.FONT_HERSHEY_DUPLEX, 0.5, (255, 255, 255))

                    if ann[-2] >= 0.5:
                        is_smile = True
                        smile_text = "smile: {:.2f}".format(ann[-2] * 100)
                        cv2.putText(img, smile_text, (x1 + 5, y1 + 30),
                                    cv2.FONT_HERSHEY_DUPLEX, 0.5,
                                    (255, 255, 255))
                        cv2.rectangle(img, (x1, y1), (x2, y2), (0, 0, 255), 2)
                    else:
                        cv2.rectangle(img, (x1, y1), (x2, y2), (0, 255, 0), 2)
            if is_smile:
                dst_file_path = file_path + '/smile_face/' + file_name
            elif has_face:
                dst_file_path = file_path + '/face/' + file_name
            else:
                dst_file_path = file_path + '/no_face/' + file_name
            cv2.imwrite(dst_file_path, img)
            print(dst_file_path)

    else:
        cam = cv2.VideoCapture('./data/linda_umu.mp4')
        # cam.set(cv2.CAP_PROP_FRAME_WIDTH, 640)
        # cam.set(cv2.CAP_PROP_FRAME_HEIGHT, 480)
        resize = FLAGS.down_scale_factor
        frame_height = cam.get(cv2.CAP_PROP_FRAME_HEIGHT) * resize
        frame_width = cam.get(cv2.CAP_PROP_FRAME_WIDTH) * resize

        max_steps = max(cfg['steps'])
        img_pad_h = max_steps - frame_height % max_steps if frame_height % max_steps > 0 else 0
        img_pad_w = max_steps - frame_width % max_steps if frame_width % max_steps > 0 else 0
        priors = prior_box_tf(
            (frame_height + img_pad_h, frame_width + img_pad_w),
            cfg['min_sizes'], cfg['steps'], cfg['clip'])

        frame_index = 0
        outputs = []
        start_time = time.time()
        while cam.isOpened():
            _, frame = cam.read()
            if frame is None:
                print('no cam')
                break
            if frame_index < 5:
                frame_index += 1
                # continue
            else:
                frame_index = 0

                img = np.float32(frame.copy())
                if resize < 1:
                    img = cv2.resize(img, (0, 0),
                                     fx=resize,
                                     fy=resize,
                                     interpolation=cv2.INTER_LINEAR)
                img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

                # pad input image to avoid unmatched shape problem
                img, pad_params = pad_input_image(img, max_steps=max_steps)

                # run model
                outputs = model(img[np.newaxis, ...])

                preds = tf.concat([
                    outputs[0][0], outputs[1][0, :, 1][..., tf.newaxis],
                    outputs[2][0, :, 1][..., tf.newaxis]
                ], -1)

                decode_preds = decode_tf(preds, priors, cfg['variances'])

                selected_indices = tf.image.non_max_suppression(
                    boxes=decode_preds[:, :4],
                    scores=decode_preds[:, -1],
                    max_output_size=tf.shape(decode_preds)[0],
                    iou_threshold=FLAGS.iou_th,
                    score_threshold=FLAGS.score_th)

                outputs = tf.gather(decode_preds, selected_indices).numpy()

                # recover padding effect
                outputs = recover_pad_output(outputs,
                                             pad_params,
                                             resize=resize)

                # calculate fps
                fps_str = "FPS: %.2f" % (1 / (time.time() - start_time))
                start_time = time.time()
                cv2.putText(frame, fps_str, (25, 50), cv2.FONT_HERSHEY_DUPLEX,
                            0.75, (0, 0, 255), 2)

            # draw results
            for prior_index in range(len(outputs)):
                draw_bbox_landm(frame, outputs[prior_index], frame_height,
                                frame_width)

            # calculate fps
            # fps_str = "FPS: %.2f" % (1 / (time.time() - start_time))
            # start_time = time.time()
            # cv2.putText(frame, fps_str, (25, 25),
            #             cv2.FONT_HERSHEY_DUPLEX, 0.75, (0, 255, 0), 2)

            # show frame
            cv2.imshow('frame', frame)
            if cv2.waitKey(1) == ord('q'):
                exit()
Exemple #14
0
def main(_argv):
    # init
    os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
    os.environ['CUDA_VISIBLE_DEVICES'] = FLAGS.gpu

    logger = tf.get_logger()
    logger.disabled = True
    logger.setLevel(logging.FATAL)
    set_memory_growth()

    cfg = load_yaml(FLAGS.cfg_path)

    # define network
    model = RetinaFaceModel(cfg,
                            training=False,
                            iou_th=FLAGS.iou_th,
                            score_th=FLAGS.score_th)

    # load checkpoint
    checkpoint_dir = './checkpoints/' + cfg['sub_name']
    checkpoint = tf.train.Checkpoint(model=model)
    if tf.train.latest_checkpoint(checkpoint_dir):
        checkpoint.restore(tf.train.latest_checkpoint(checkpoint_dir))
        print("[*] load ckpt from {}.".format(
            tf.train.latest_checkpoint(checkpoint_dir)))
    else:
        print("[*] Cannot find ckpt from {}.".format(checkpoint_dir))
        exit()

    if not FLAGS.webcam:
        if not os.path.exists(FLAGS.img_path):
            print(f"cannot find image path from {FLAGS.img_path}")
            exit()

        print("[*] Processing on single image {}".format(FLAGS.img_path))

        img_raw = cv2.imread(FLAGS.img_path)
        img_height_raw, img_width_raw, _ = img_raw.shape
        img = np.float32(img_raw.copy())

        if FLAGS.down_scale_factor < 1.0:
            img = cv2.resize(img, (0, 0),
                             fx=FLAGS.down_scale_factor,
                             fy=FLAGS.down_scale_factor,
                             interpolation=cv2.INTER_LINEAR)
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

        # pad input image to avoid unmatched shape problem
        img, pad_params = pad_input_image(img, max_steps=max(cfg['steps']))

        # run model
        outputs = model(img[np.newaxis, ...]).numpy()

        # recover padding effect
        outputs = recover_pad_output(outputs, pad_params)

        # draw and save results
        save_img_path = os.path.join('out_' + os.path.basename(FLAGS.img_path))
        for prior_index in range(len(outputs)):
            draw_bbox_landm(img_raw, outputs[prior_index], img_height_raw,
                            img_width_raw)
            cv2.imwrite(save_img_path, img_raw)
        print(f"[*] save result at {save_img_path}")

    else:
        cam = cv2.VideoCapture("./videos/Bentall_Centra.MP4")

        start_time = time.time()
        while True:
            _, frame = cam.read()
            if frame is None:
                print("no cam input")

            frame_height, frame_width, _ = frame.shape
            img = np.float32(frame.copy())
            if FLAGS.down_scale_factor < 1.0:
                img = cv2.resize(img, (0, 0),
                                 fx=FLAGS.down_scale_factor,
                                 fy=FLAGS.down_scale_factor,
                                 interpolation=cv2.INTER_LINEAR)
            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

            # pad input image to avoid unmatched shape problem
            img, pad_params = pad_input_image(img, max_steps=max(cfg['steps']))

            # run model
            outputs = model(img[np.newaxis, ...]).numpy()

            # recover padding effect
            outputs = recover_pad_output(outputs, pad_params)

            # draw results
            for prior_index in range(len(outputs)):
                draw_bbox_landm(frame, outputs[prior_index], frame_height,
                                frame_width)

            # calculate fps
            fps_str = "FPS: %.2f" % (1 / (time.time() - start_time))
            start_time = time.time()
            cv2.putText(frame, fps_str, (25, 25), cv2.FONT_HERSHEY_DUPLEX,
                        0.75, (0, 255, 0), 2)

            # show frame
            cv2.imshow('frame', frame)
            if cv2.waitKey(1) == ord('q'):
                exit()
Exemple #15
0
def main(_):
    # init
    os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
    os.environ['CUDA_VISIBLE_DEVICES'] = FLAGS.gpu

    logger = tf.get_logger()
    logger.disabled = True
    logger.setLevel(logging.FATAL)
    set_memory_growth()

    cfg = load_yaml(FLAGS.cfg_path)

    # define network
    generator = RRDB_Model(cfg['input_size'], cfg['ch_size'], cfg['network_G'])
    generator.summary(line_length=80)
    discriminator = DiscriminatorVGG128(cfg['gt_size'], cfg['ch_size'])
    discriminator.summary(line_length=80)

    # load dataset
    train_dataset = load_dataset(cfg, 'train_dataset', shuffle=False)

    # define optimizer
    learning_rate_G = MultiStepLR(cfg['lr_G'], cfg['lr_steps'], cfg['lr_rate'])
    learning_rate_D = MultiStepLR(cfg['lr_D'], cfg['lr_steps'], cfg['lr_rate'])
    optimizer_G = tf.keras.optimizers.Adam(learning_rate=learning_rate_G,
                                           beta_1=cfg['adam_beta1_G'],
                                           beta_2=cfg['adam_beta2_G'])
    optimizer_D = tf.keras.optimizers.Adam(learning_rate=learning_rate_D,
                                           beta_1=cfg['adam_beta1_D'],
                                           beta_2=cfg['adam_beta2_D'])

    # define losses function
    pixel_loss_fn = PixelLoss(criterion=cfg['pixel_criterion'])
    fea_loss_fn = ContentLoss(criterion=cfg['feature_criterion'])
    gen_loss_fn = GeneratorLoss(gan_type=cfg['gan_type'])
    dis_loss_fn = DiscriminatorLoss(gan_type=cfg['gan_type'])

    # load checkpoint
    checkpoint_dir = './checkpoints/' + cfg['sub_name']
    checkpoint = tf.train.Checkpoint(step=tf.Variable(0, name='step'),
                                     optimizer_G=optimizer_G,
                                     optimizer_D=optimizer_D,
                                     model=generator,
                                     discriminator=discriminator)
    manager = tf.train.CheckpointManager(checkpoint=checkpoint,
                                         directory=checkpoint_dir,
                                         max_to_keep=3)
    if manager.latest_checkpoint:
        checkpoint.restore(manager.latest_checkpoint)
        print('[*] load ckpt from {} at step {}.'.format(
            manager.latest_checkpoint, checkpoint.step.numpy()))
    else:
        if cfg['pretrain_name'] is not None:
            pretrain_dir = './checkpoints/' + cfg['pretrain_name']
            if tf.train.latest_checkpoint(pretrain_dir):
                checkpoint.restore(tf.train.latest_checkpoint(pretrain_dir))
                checkpoint.step.assign(0)
                print("[*] training from pretrain model {}.".format(
                    pretrain_dir))
            else:
                print(
                    "[*] cannot find pretrain model {}.".format(pretrain_dir))
        else:
            print("[*] training from scratch.")

    # define training step function
    @tf.function
    def train_step(lr, hr):
        with tf.GradientTape(persistent=True) as tape:
            sr = generator(lr, training=True)
            hr_output = discriminator(hr, training=True)
            sr_output = discriminator(sr, training=True)

            losses_G = {}
            losses_D = {}
            losses_G['reg'] = tf.reduce_sum(generator.losses)
            losses_D['reg'] = tf.reduce_sum(discriminator.losses)
            losses_G['pixel'] = cfg['w_pixel'] * pixel_loss_fn(hr, sr)
            losses_G['feature'] = cfg['w_feature'] * fea_loss_fn(hr, sr)
            losses_G['gan'] = cfg['w_gan'] * gen_loss_fn(hr_output, sr_output)
            losses_D['gan'] = dis_loss_fn(hr_output, sr_output)
            total_loss_G = tf.add_n([l for l in losses_G.values()])
            total_loss_D = tf.add_n([l for l in losses_D.values()])

        grads_G = tape.gradient(total_loss_G, generator.trainable_variables)
        grads_D = tape.gradient(total_loss_D,
                                discriminator.trainable_variables)
        optimizer_G.apply_gradients(zip(grads_G,
                                        generator.trainable_variables))
        optimizer_D.apply_gradients(
            zip(grads_D, discriminator.trainable_variables))

        return total_loss_G, total_loss_D, losses_G, losses_D

    # training loop
    summary_writer = tf.summary.create_file_writer('./logs/' + cfg['sub_name'])
    prog_bar = ProgressBar(cfg['niter'], checkpoint.step.numpy())
    remain_steps = max(cfg['niter'] - checkpoint.step.numpy(), 0)

    for lr, hr in train_dataset.take(remain_steps):
        checkpoint.step.assign_add(1)
        steps = checkpoint.step.numpy()

        total_loss_G, total_loss_D, losses_G, losses_D = train_step(lr, hr)

        prog_bar.update(
            "loss_G={:.4f}, loss_D={:.4f}, lr_G={:.1e}, lr_D={:.1e}".format(
                total_loss_G.numpy(), total_loss_D.numpy(),
                optimizer_G.lr(steps).numpy(),
                optimizer_D.lr(steps).numpy()))

        if steps % 10 == 0:
            with summary_writer.as_default():
                tf.summary.scalar('loss_G/total_loss',
                                  total_loss_G,
                                  step=steps)
                tf.summary.scalar('loss_D/total_loss',
                                  total_loss_D,
                                  step=steps)
                for k, l in losses_G.items():
                    tf.summary.scalar('loss_G/{}'.format(k), l, step=steps)
                for k, l in losses_D.items():
                    tf.summary.scalar('loss_D/{}'.format(k), l, step=steps)

                tf.summary.scalar('learning_rate_G',
                                  optimizer_G.lr(steps),
                                  step=steps)
                tf.summary.scalar('learning_rate_D',
                                  optimizer_D.lr(steps),
                                  step=steps)

        if steps % cfg['save_steps'] == 0:
            manager.save()
            print("\n[*] save ckpt file at {}".format(
                manager.latest_checkpoint))

    print("\n [*] training done!")
def main(_argv):
    # init
    os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
    os.environ['CUDA_VISIBLE_DEVICES'] = FLAGS.gpu

    logger = tf.get_logger()
    logger.disabled = True
    logger.setLevel(logging.FATAL)
    set_memory_growth()

    cfg = load_yaml(FLAGS.cfg_path)
    aligner = FaceAligner()
    # define network
    model = RetinaFaceModel(cfg,
                            training=False,
                            iou_th=FLAGS.iou_th,
                            score_th=FLAGS.score_th)

    # load checkpoint
    checkpoint_dir = './checkpoints/' + cfg['sub_name']
    checkpoint = tf.train.Checkpoint(model=model)
    if tf.train.latest_checkpoint(checkpoint_dir):
        checkpoint.restore(tf.train.latest_checkpoint(checkpoint_dir))
        print("[*] load ckpt from {}.".format(
            tf.train.latest_checkpoint(checkpoint_dir)))
    else:
        print("[*] Cannot find ckpt from {}.".format(checkpoint_dir))
        exit()
    if FLAGS.input_stream == '0':
        input_stream = 0
    elif FLAGS.input_stream == 'rtsp':
        input_stream = settings.RTSP_ADDR
    else:
        input_stream = FLAGS.input_stream
    cam = cv2.VideoCapture(input_stream)  #("/home/hao/Videos/Webcam/3.webm")
    mbv2 = tf.keras.models.load_model(settings.CHECKPOINT_PATH)
    anchor_dataset = np.load(settings.ANCHOR_PATH)['arr_0']
    label_dataset = np.load(settings.LABEL_PATH)['arr_0']
    start_time = time.time()
    i = 0
    while cam.isOpened():
        _, frame = cam.read()
        if frame is None:
            print("no cam input")

        frame_height, frame_width, _ = frame.shape
        img = np.float32(frame.copy())
        if FLAGS.down_scale_factor < 1.0:
            img = cv2.resize(img, (0, 0),
                             fx=FLAGS.down_scale_factor,
                             fy=FLAGS.down_scale_factor,
                             interpolation=cv2.INTER_LINEAR)
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

        # pad input image to avoid unmatched shape problem
        img, pad_params = pad_input_image(img, max_steps=max(cfg['steps']))

        # run model
        outputs = model(img[np.newaxis, ...]).numpy()

        # recover padding effect
        outputs = recover_pad_output(outputs, pad_params)

        # draw results
        for prior_index in range(len(outputs)):
            ann = outputs[prior_index]
            b_box = int(ann[0] * frame_width), int(ann[1] * frame_height), \
                     int(ann[2] * frame_width), int(ann[3] * frame_height)
            if (b_box[0] < 0) or (b_box[1] < 0) or (
                    b_box[2] >= frame_width) or (b_box[3] >= frame_height):
                continue
            keypoints = {
                'left_eye': (ann[4] * frame_width, ann[5] * frame_height),
                'right_eye': (ann[6] * frame_width, ann[7] * frame_height),
                'nose': (ann[8], ann[9]),
                'left_mouth': (ann[10] * frame_width, ann[11] * frame_height),
                'right_mouth': (ann[12] * frame_width, ann[13] * frame_height),
            }
            out_frame = aligner.align(frame, keypoints, b_box)
            scaled = out_frame  #cv2.resize(out_frame, (settings.IMAGE_SIZE, settings.IMAGE_SIZE), interpolation=cv2.INTER_CUBIC)
            scaled_reshape = scaled.reshape(-1, settings.IMAGE_SIZE,
                                            settings.IMAGE_SIZE, 3)
            embed_vector = mbv2(scaled_reshape / 255.0)
            label, prob = classify(embed_vector, anchor_dataset, label_dataset)
            if prob < 0.5:
                label = "Unknown"
            cv2.rectangle(frame, (b_box[0], b_box[1]), (b_box[2], b_box[3]),
                          (0, 255, 0), 2)
            cv2.putText(frame, label, (b_box[0], b_box[1]),
                        cv2.FONT_HERSHEY_DUPLEX, 0.5, (255, 255, 255))

            text = "{:.4f}".format(prob)
            cv2.putText(frame, text, (b_box[0], b_box[1] + 15),
                        cv2.FONT_HERSHEY_DUPLEX, 0.5, (255, 255, 255))

        # calculate fps
        fps_str = "FPS: %.2f" % (1 / (time.time() - start_time))
        start_time = time.time()
        cv2.putText(frame, fps_str, (25, 25), cv2.FONT_HERSHEY_DUPLEX, 0.75,
                    (0, 255, 0), 2)
        i += 1
        # show frame
        # cv2.imwrite('UNKNOWN/4/'+str(i)+'.jpeg', frame)
        cv2.imshow("frame", frame)
        if cv2.waitKey(1) == ord('q'):
            exit()
Exemple #17
0
def main(_):
    # init
    os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
    os.environ['CUDA_VISIBLE_DEVICES'] = FLAGS.gpu
    set_memory_growth()

    logger = tf.get_logger()
    logger.disabled = True
    logger.setLevel(logging.FATAL)

    cfg = load_yaml(FLAGS.cfg_path)
    uv_weight_mask = cv2.imread(cfg['uv_weight_mask']) / 255.

    # define network
    generator = PRNet_Model(cfg['input_size'], cfg['ch_size'])
    generator.summary(line_length=80)

    # load dataset
    train_dataset = load_dataset(cfg,
                                 shuffle=True,
                                 num_workers=cfg['num_workers'])

    # define optimizer
    learning_rate_G = MultiStepLR(cfg['lr_G'], cfg['lr_steps'], cfg['lr_rate'])
    optimizer_G = tf.keras.optimizers.Adam(learning_rate=learning_rate_G,
                                           beta_1=cfg['adam_beta1_G'],
                                           beta_2=cfg['adam_beta2_G'])

    # define losses function
    loss_fn = WeightedMSE(uv_weight_mask)

    # load checkpoint
    checkpoint_dir = './checkpoints/' + cfg['sub_name']
    checkpoint = tf.train.Checkpoint(step=tf.Variable(0, name='step'),
                                     optimizer_G=optimizer_G,
                                     model=generator)
    manager = tf.train.CheckpointManager(checkpoint=checkpoint,
                                         directory=checkpoint_dir,
                                         max_to_keep=3)
    if manager.latest_checkpoint:
        checkpoint.restore(manager.latest_checkpoint)
        print('[*] load ckpt from {} at step {}.'.format(
            manager.latest_checkpoint, checkpoint.step.numpy()))
    else:
        if cfg['pretrain_name'] is not None:
            pretrain_dir = './checkpoints/' + cfg['pretrain_name']
            if tf.train.latest_checkpoint(pretrain_dir):
                checkpoint.restore(tf.train.latest_checkpoint(pretrain_dir))
                checkpoint.step.assign(0)
                print("[*] training from pretrain model {}.".format(
                    pretrain_dir))
            else:
                print(
                    "[*] cannot find pretrain model {}.".format(pretrain_dir))
        else:
            print("[*] training from scratch.")

    # define training step function
    @tf.function
    def train_step(img, pos):
        with tf.GradientTape() as tape_G:
            pre = generator(img, training=True)

            losses_G = {}
            losses_G['pixel'] = loss_fn(pos, pre)
            total_loss_G = tf.add_n([l for l in losses_G.values()])

        grads_G = tape_G.gradient(total_loss_G, generator.trainable_variables)
        optimizer_G.apply_gradients(zip(grads_G,
                                        generator.trainable_variables))

        return total_loss_G, losses_G

    # training loop
    summary_writer = tf.summary.create_file_writer('./logs/' + cfg['sub_name'])
    niter = int(cfg['train_dataset']['num_samples'] * cfg['epoch'] /
                cfg['batch_size'])
    prog_bar = ProgressBar(niter, checkpoint.step.numpy())
    remain_steps = max(niter - checkpoint.step.numpy(), 0)

    for sample in take(remain_steps, train_dataset):
        checkpoint.step.assign_add(1)
        steps = checkpoint.step.numpy()
        img, pos = sample['Image'], sample['Posmap']

        total_loss_G, losses_G = train_step(img, pos)

        prog_bar.update_gan(total_loss_G.numpy(),
                            optimizer_G.lr(steps).numpy())

        if steps % cfg['log_steps'] == 0:
            with summary_writer.as_default():
                tf.summary.scalar('loss_G/total_loss',
                                  total_loss_G,
                                  step=steps)
                for k, l in losses_G.items():
                    tf.summary.scalar('loss_G/{}'.format(k), l, step=steps)

                tf.summary.scalar('learning_rate_G',
                                  optimizer_G.lr(steps),
                                  step=steps)

        if steps % cfg['save_steps'] == 0:
            manager.save()
            print('\n[*] save ckpt file at {}'.format(
                manager.latest_checkpoint))

    print("\n[*] training done!")
def main(_):
    set_memory_growth()

    cfg = load_yaml(FLAGS.cfg_path)
    model = ArcFaceModel(size=cfg['input_size'],
                         backbone_type=cfg['backbone_type'],
                         num_classes=cfg['num_classes'],
                         head_type=cfg['head_type'],
                         embd_shape=cfg['embd_shape'],
                         w_decay=cfg['w_decay'],
                         training=True)
    learning_rate = tf.constant(cfg['base_lr'])
    optimizer = tf.keras.optimizers.SGD(learning_rate=learning_rate,
                                        momentum=0.9,
                                        nesterov=True)
    loss_fn = SoftmaxLoss()

    ckpt_path = tf.train.latest_checkpoint('./checkpoints/train_' +
                                           cfg['sub_name'])
    if ckpt_path is not None:
        print("[*] load ckpt from {}".format(ckpt_path))
        model.load_weights(ckpt_path)
    else:
        print("[*] training from scratch.")
    model.compile(optimizer=optimizer, loss=loss_fn)

    # resnet_model = tf.keras.Model(inputs=model.get_layer('resnet50').input,
    #                                 outputs=model.get_layer('resnet50').output)
    # resnet_head = tf.keras.Model(inputs=resnet_model.input,
    #                                 outputs=resnet_model.get_layer('conv2_block1_add').output)
    # resnet_tail = split(resnet_model, 18, 1000) # conv2_block1_out
    # output_model = tf.keras.Model(inputs=model.get_layer('OutputLayer').input,
    #                                 outputs=model.get_layer('OutputLayer').output)
    # archead = tf.keras.Model(inputs=model.get_layer('ArcHead').input,
    #                                 outputs=model.get_layer('ArcHead').output)

    temp1 = np.ones((62, 112, 3))
    temp2 = np.zeros((50, 112, 3))
    masked_img = np.concatenate([temp1, temp2], axis=0).reshape(1, 112, 112, 3)

    temp1 = np.ones((1, 18, 28, 256))
    temp2 = np.zeros((1, 10, 28, 256))
    masked = np.concatenate([temp1, temp2], axis=1)
    # inputs = Input((112, 112, 3))
    # labels = Input([])
    # s = resnet_head(inputs)
    # s = tf.keras.layers.Multiply()([s, tf.constant(masked, dtype=tf.float32)])
    # s = resnet_tail(s)
    # s = output_model(s)
    # s = archead([s, labels])
    # new_model = Model((inputs, labels), s)
    # new_model.summary()

    # new_model.compile(optimizer=optimizer, loss=loss_fn)

    path_to_data = '/home/anhdq23/Desktop/nguyen/data/AR/test2'
    anchor_names = get_gallery_pr2(path_to_data)  # From 1 to 100
    name_dicts = get_probe_pr2(
        path_to_data)  # Dictionary: {anchor_name:[name_image, ...]}

    augment = ImgAugTransform()
    import faiss

    if FLAGS.mode == 'eager_tf':
        top_left_all = [0.012]
        best_acc = 0.8
        for epochs in range(cfg['epochs']):
            logging.info("Shuffle ms1m dataset.")
            dataset_len = cfg['num_samples']
            steps_per_epoch = dataset_len // cfg['batch_size']
            train_dataset = dataset.load_tfrecord_dataset(
                cfg['train_dataset'],
                cfg['batch_size'],
                cfg['binary_img'],
                is_ccrop=cfg['is_ccrop'])

            sub_train_dataset = dataset.load_tfrecord_dataset(
                cfg['train_dataset'],
                cfg['batch_size'],
                cfg['binary_img'],
                is_ccrop=cfg['is_ccrop'])

            for batch, ((x, y), (sub_x, sub_y)) in enumerate(
                    zip(train_dataset, sub_train_dataset)):
                x0_new = np.array(x[0], dtype=np.uint8)
                x1_new = np.array(x[1], dtype=np.float32)
                for i in np.arange(len(x0_new)):
                    x0_new[i] = augment(x0_new[i])
                temp = np.array(x0_new, dtype=np.float32) / 255.0
                temp = np.multiply(temp, masked_img)

                sub_x0_new = np.array(sub_x[0], dtype=np.uint8)
                sub_x1_new = np.array(sub_x[1], dtype=np.float32)
                for i in np.arange(len(sub_x0_new)):
                    sub_x0_new[i] = augment(sub_x0_new[i])
                sub_temp = np.array(sub_x0_new, dtype=np.float32) / 255.0

                model.train_on_batch(*((sub_temp, sub_x1_new), sub_x1_new))
                loss = model.train_on_batch(*((temp, x1_new), x1_new))

                if batch % 50 == 0:
                    verb_str = "Epoch {}/{}: {}/{}, loss={:.6f}, lr={:.6f}"
                    print(
                        verb_str.format(
                            epochs, cfg['epochs'], batch, steps_per_epoch,
                            loss, cfg['base_lr'] / (1.0 + cfg['w_decay'] *
                                                    (epochs * 45489 + batch))))

                    if batch % cfg['save_steps'] == 0:
                        resnet_model = tf.keras.Model(
                            inputs=model.get_layer('resnet50').input,
                            outputs=model.get_layer('resnet50').output)

                        output_model = tf.keras.Model(
                            inputs=model.get_layer('OutputLayer').input,
                            outputs=model.get_layer('OutputLayer').output)

                        database_image_names = []
                        database_feature_list = []
                        for anchor_name in anchor_names:
                            img1 = Image.open(
                                os.path.join(path_to_data, anchor_name))
                            img1 = img1.resize((112, 112))
                            img1 = np.array(img1) / 255.0
                            img1 = np.multiply(img1, masked_img)

                            fc1 = resnet_model.predict(
                                img1.reshape((1, 112, 112, 3)))
                            fc1 = output_model.predict(fc1)
                            norm_fc1 = preprocessing.normalize(fc1.reshape(
                                (1, 512)),
                                                               norm='l2',
                                                               axis=1)
                            database_image_names.append(anchor_name)
                            database_feature_list.append(norm_fc1)

                        index_flat = faiss.IndexFlatL2(512)
                        gpu_index_flat = index_flat
                        gpu_index_flat.add(
                            np.array(database_feature_list).reshape(
                                (-1, 512)))  # add vectors to the index
                        count = 0
                        for key in list(name_dicts.keys()):
                            for name in name_dicts[key]:
                                img2 = Image.open(
                                    os.path.join(path_to_data, name)).resize(
                                        (112, 112))
                                img2 = img2.resize((112, 112))
                                img2 = np.array(img2) / 255.0
                                img2 = np.multiply(img2, masked_img)

                                fc2 = resnet_model.predict(
                                    img2.reshape((1, 112, 112, 3)))
                                fc2 = output_model.predict(fc2)
                                norm_fc2 = preprocessing.normalize(fc2.reshape(
                                    (1, 512)),
                                                                   norm='l2',
                                                                   axis=1)

                                D, I = gpu_index_flat.search(
                                    norm_fc2, k=1)  # actual search
                                if name[0:5] == database_image_names[I[0]
                                                                     [0]][0:5]:
                                    count += 1
                        acc = count / 600.0
                        if acc > best_acc:
                            best_acc = acc
                            print('[*] save ckpt file!')
                            model.save_weights(
                                'checkpoints/{}/e_{}_b_{}.ckpt'.format(
                                    cfg['sub_name'], epochs,
                                    batch % steps_per_epoch))
                        print("The current acc: {:.6f} ".format(acc))
                        print("The best acc: {:.6f} ".format(best_acc))

                    model.save_weights('checkpoints/train_{}/{}.ckpt'.format(
                        cfg['sub_name'], cfg['sub_name']))
Exemple #19
0
def main(_argv):
    # init

    face_aligner = FaceAligner()

    os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
    os.environ['CUDA_VISIBLE_DEVICES'] = FLAGS.gpu

    logger = tf.get_logger()
    logger.disabled = True
    logger.setLevel(logging.FATAL)
    set_memory_growth()

    cfg = load_yaml(FLAGS.cfg_path)

    # define network
    model = RetinaFaceModel(cfg,
                            training=False,
                            iou_th=FLAGS.iou_th,
                            score_th=FLAGS.score_th)

    # load checkpoint
    checkpoint_dir = './checkpoints/' + cfg['sub_name']
    checkpoint = tf.train.Checkpoint(model=model)
    if tf.train.latest_checkpoint(checkpoint_dir):
        checkpoint.restore(tf.train.latest_checkpoint(checkpoint_dir))
        print("[*] load ckpt from {}.".format(
            tf.train.latest_checkpoint(checkpoint_dir)))
    else:
        print("[*] Cannot find ckpt from {}.".format(checkpoint_dir))
        exit()

    if not FLAGS.webcam:
        save_count = 0
        for path, subdirs, files in os.walk(FLAGS.img_path):

            for name in files:
                if name.endswith('.jpg'):
                    img_path = os.path.join(path, name)

                    if not os.path.exists(img_path):
                        print(f"cannot find image path from {img_path}")
                        exit()

                    if save_count < FLAGS.img_num:
                        print("[*] Processing on single image {}".format(
                            img_path))

                        img_raw = cv2.imread(img_path)
                        img_height_raw, img_width_raw, _ = img_raw.shape
                        img = np.float32(img_raw.copy())

                        if FLAGS.down_scale_factor < 1.0:
                            img = cv2.resize(img, (0, 0),
                                             fx=FLAGS.down_scale_factor,
                                             fy=FLAGS.down_scale_factor,
                                             interpolation=cv2.INTER_LINEAR)
                        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

                        # pad input image to avoid unmatched shape problem
                        img, pad_params = pad_input_image(img,
                                                          max_steps=max(
                                                              cfg['steps']))

                        # run model
                        outputs = model(img[np.newaxis, ...]).numpy()

                        # recover padding effect
                        outputs = recover_pad_output(outputs, pad_params)

                        # draw and save results
                        save_img_path = os.path.join(
                            'out_' + os.path.basename(img_path))

                        for prior_index in range(len(outputs)):
                            draw_bbox_landm(img_raw, outputs[prior_index],
                                            img_height_raw, img_width_raw)
                            cv2.imwrite(save_img_path, img_raw)
                        print(f"[*] save result at {save_img_path}")
                        save_count += 1

    else:
        cam = cv2.VideoCapture(0)

        start_time = time.time()

        while True:
            _, frame = cam.read()
            if frame is None:
                print("no cam input")

            frame_height, frame_width, _ = frame.shape

            orig_frame = frame.copy()

            face = None

            img = cv2.resize(frame, (512, 512))
            img = np.float32(frame.copy())
            if FLAGS.down_scale_factor < 1.0:
                img = cv2.resize(img, (0, 0),
                                 fx=FLAGS.down_scale_factor,
                                 fy=FLAGS.down_scale_factor,
                                 interpolation=cv2.INTER_LINEAR)

            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

            # pad input image to avoid unmatched shape problem
            img, pad_params = pad_input_image(img, max_steps=max(cfg['steps']))

            # run model
            start_time = time.time()
            outputs = model(img[np.newaxis, ...]).numpy()
            inference_time = f"Inf: {time.time() - start_time}"

            cv2.putText(frame, inference_time, (25, 50),
                        cv2.FONT_HERSHEY_DUPLEX, 0.75, (0, 255, 0), 2)

            # recover padding effect
            outputs = recover_pad_output(outputs, pad_params)

            # draw results
            for prior_index in range(len(outputs)):
                preds = decode_predictions((frame_width, frame_height),
                                           outputs)

                for key, value in preds.items():
                    bbox = value[0]['bbox']
                    left_eye = value[0]['left_eye']
                    right_eye = value[0]['right_eye']

                    # Our face ROI
                    face = orig_frame[bbox[1]:bbox[3], bbox[0]:bbox[2]]

                    # Eyes
                    x1_le = left_eye[0] - 25
                    y1_le = left_eye[1] - 25
                    x2_le = left_eye[0] + 25
                    y2_le = left_eye[1] + 25

                    x1_re = right_eye[0] - 25
                    y1_re = right_eye[1] - 25
                    x2_re = right_eye[0] + 25
                    y2_re = right_eye[1] + 25

                    if left_eye[1] > right_eye[1]:
                        A = (right_eye[0], left_eye[1])
                    else:
                        A = (left_eye[0], right_eye[1])

                    # Calc our rotating degree
                    delta_x = right_eye[0] - left_eye[0]
                    delta_y = right_eye[1] - left_eye[1]
                    angle = np.arctan(
                        delta_y / (delta_x + 1e-17))  # avoid devision by zero
                    angle = (angle * 180) / np.pi

                    # compute the desired right eye x-coordinate based on the
                    # desired x-coordinate of the left eye
                    desiredRightEyeX = 1.0 - 0.35

                    # determine the scale of the new resulting image by taking
                    # the ratio of the distance between eyes in the *current*
                    # image to the ratio of distance between eyes in the
                    # *desired* image
                    dist = np.sqrt((delta_x**2) + (delta_y**2))
                    desiredDist = (desiredRightEyeX - 0.35)
                    desiredDist *= 256
                    scale = desiredDist / dist

                    eyesCenter = ((left_eye[0] + right_eye[0]) // 2,
                                  (left_eye[1] + right_eye[1]) // 2)

                    cv2.circle(frame, A, 5, (255, 0, 0), -1)

                    cv2.putText(frame, str(int(angle)), (x1_le - 15, y1_le),
                                cv2.FONT_HERSHEY_DUPLEX, 0.75, (0, 255, 0), 2)

                    cv2.line(frame, right_eye, left_eye, (0, 200, 200), 3)
                    cv2.line(frame, left_eye, A, (0, 200, 200), 3)
                    cv2.line(frame, right_eye, A, (0, 200, 200), 3)

                    cv2.line(frame, (left_eye[0], left_eye[1]),
                             (right_eye[0], right_eye[1]), (0, 200, 200), 3)

                    rotated = face_aligner.align(orig_frame, left_eye,
                                                 right_eye)

                draw_bbox_landm(frame, outputs[prior_index], frame_height,
                                frame_width)

            # calculate fps
            fps_str = "FPS: %.2f" % (1 / (time.time() - start_time))
            start_time = time.time()
            cv2.putText(frame, fps_str, (25, 25), cv2.FONT_HERSHEY_DUPLEX,
                        0.75, (0, 255, 0), 2)

            # show frame
            cv2.imshow('frame', frame)
            if face is not None:
                cv2.imshow('face aligned', rotated)

            if cv2.waitKey(1) == ord('q'):
                exit()
Exemple #20
0
def main(_argv):
    os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
    os.environ['CUDA_VISIBLE_DEVICES'] = FLAGS.gpu

    logger = tf.get_logger()
    logger.disabled = True
    logger.setLevel(logging.FATAL)
    set_memory_growth()

    cfg = load_yaml(FLAGS.cfg_path)

    model = ArcFaceModel(size=cfg['input_size'],
                         backbone_type=cfg['backbone_type'],
                         training=False)

    ckpt_path = tf.train.latest_checkpoint('./checkpoints/' + cfg['sub_name'])
    print(ckpt_path)
    if ckpt_path is not None:
        print("[*] load ckpt from {}".format(ckpt_path))
        model.load_weights(ckpt_path)
    else:
        print("[*] Cannot find ckpt from {}.".format(ckpt_path))
        exit()

    if FLAGS.img_path:
        print("[*] Encode {} to ./output_embeds.npy".format(FLAGS.img_path))
        img = cv2.imread(FLAGS.img_path)
        img = cv2.resize(img, (cfg['input_size'], cfg['input_size']))
        img = img.astype(np.float32) / 255.
        if len(img.shape) == 3:
            img = np.expand_dims(img, 0)
        embeds = l2_norm(model(img))
        print(embeds.shape)
        np.save('./output_embeds.npy', embeds)
    else:
        print("[*] Loading LFW, AgeDB30 and CFP-FP...")
        lfw, agedb_30, cfp_fp, lfw_issame, agedb_30_issame, cfp_fp_issame = \
            get_val_data(cfg['test_dataset'])

        print("[*] Perform Evaluation on LFW...")
        acc_lfw, best_th = perform_val(cfg['embd_shape'],
                                       cfg['batch_size'],
                                       model,
                                       lfw,
                                       lfw_issame,
                                       is_ccrop=cfg['is_ccrop'])
        print("    acc {:.4f}, th: {:.2f}".format(acc_lfw, best_th))

        print("[*] Perform Evaluation on AgeDB30...")
        acc_agedb30, best_th = perform_val(cfg['embd_shape'],
                                           cfg['batch_size'],
                                           model,
                                           agedb_30,
                                           agedb_30_issame,
                                           is_ccrop=cfg['is_ccrop'])
        print("    acc {:.4f}, th: {:.2f}".format(acc_agedb30, best_th))

        print("[*] Perform Evaluation on CFP-FP...")
        acc_cfp_fp, best_th = perform_val(cfg['embd_shape'],
                                          cfg['batch_size'],
                                          model,
                                          cfp_fp,
                                          cfp_fp_issame,
                                          is_ccrop=cfg['is_ccrop'])
        print("    acc {:.4f}, th: {:.2f}".format(acc_cfp_fp, best_th))
Exemple #21
0
def main(_argv):
    os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
    os.environ['CUDA_VISIBLE_DEVICES'] = FLAGS.gpu

    logger = tf.get_logger()
    logger.disabled = True
    logger.setLevel(logging.FATAL)
    set_memory_growth()

    cfg = load_yaml(FLAGS.cfg_path)

    model = ArcFaceModel(size=cfg['input_size'],
                         backbone_type=cfg['backbone_type'],
                         training=False)

    ckpt_path = tf.train.latest_checkpoint('./checkpoints/' + cfg['sub_name'])
    if ckpt_path is not None:
        print("[*] load ckpt from {}".format(ckpt_path))
        model.load_weights(ckpt_path)
    else:
        print("[*] Cannot find ckpt from {}.".format(ckpt_path))
        exit()

    if FLAGS.img_path:

        print("Check Start!!!")

        file_dir = "C:/Users/smgg/Desktop/dataset/superjunior/all3.jpeg"
        npy_dir = "/SCLab/newTWICE_id/*,npy"

        img_list = glob.glob(file_dir)
        npy_list = glob.glob(npy_dir)

        for img_name in img_list:
            img = cv2.cvtColor(cv2.imread(img_name), cv2.COLOR_BGR2RGB)
            detector = MTCNN()
            data_list = detector.detect_faces(img)

            for data in data_list:
                xmin, ymin, width, height = data['box']
                xmax = xmin + width
                ymax = ymin + height

                face_image = img[ymin:ymax, xmin: xmax, :]
                face_image = cv2.cvtColor(face_image, cv2.COLOR_RGB2BGR)

                # cv2.imshow('', face_image)
                # cv2.waitKey(0)

                img_resize = cv2.resize(face_image, (cfg['input_size'], cfg['input_size']))
                img_resize = img_resize.astype(np.float32) / 255.
                if len(img_resize.shape) == 3:
                    img_resize = np.expand_dims(img_resize, 0)
                embeds = l2_norm(model(img_resize))

                i = 0
                cv2.rectangle(img, (xmin, ymin), (xmax, ymax), (0, 255, 0), 1)
                for npy_name in npy_list:

                    name_embeds = np.load(npy_name)
                    # print(1)

                    if distance(embeds, name_embeds,1) < 0.37:
                        i = i + 1
                        name = npy_name.split('/')[5].split('\\')[1].split('.npy')[0]

                        cv2.putText(img, name, (xmin, ymin - 15*(i)), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 1,cv2.LINE_AA)


                    # else:
                    #     cv2.putText(img, "Unknown", (xmin, ymin - 10), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 1, cv2.LINE_AA)
                    #     cv2.rectangle(img, (xmin, ymin), (xmax, ymax), (0, 255, 0), 1)
            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
            cv2.imshow('', img)
            cv2.waitKey(0)
Exemple #22
0
def main(_):
    '''
    Train for one epoch to get supernet , then random sample 50 architectures for finetuning.
    This structure is basically the same as train_search.py
    TODO: Add PGD here and calculate FSP
    '''
    # init
    os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
    os.environ['CUDA_VISIBLE_DEVICES'] = FLAGS.gpu

    logger = tf.get_logger()
    logger.disabled = True
    logger.setLevel(logging.FATAL)
    set_memory_growth()

    cfg = load_yaml(FLAGS.cfg_path)

    # define network
    sna = SearchNetArch(cfg)
    sna.model.summary(line_length=80)
    print("param size = {:f}MB".format(count_parameters_in_MB(sna.model)))

    # load dataset
    t_split = f"train[0%:{int(cfg['train_portion'] * 100)}%]"
    v_split = f"train[{int(cfg['train_portion'] * 100)}%:100%]"
    train_dataset = load_cifar10_dataset(
        cfg['batch_size'],
        split=t_split,
        shuffle=True,
        drop_remainder=True,
        using_normalize=cfg['using_normalize'],
        using_crop=cfg['using_crop'],
        using_flip=cfg['using_flip'],
        using_cutout=cfg['using_cutout'],
        cutout_length=cfg['cutout_length'])
    val_dataset = load_cifar10_dataset(cfg['batch_size'],
                                       split=v_split,
                                       shuffle=True,
                                       drop_remainder=True,
                                       using_normalize=cfg['using_normalize'],
                                       using_crop=cfg['using_crop'],
                                       using_flip=cfg['using_flip'],
                                       using_cutout=cfg['using_cutout'],
                                       cutout_length=cfg['cutout_length'])

    # define optimizer
    steps_per_epoch = int(cfg['dataset_len'] * cfg['train_portion'] //
                          cfg['batch_size'])
    learning_rate = CosineAnnealingLR(initial_learning_rate=cfg['init_lr'],
                                      t_period=cfg['epoch'] * steps_per_epoch,
                                      lr_min=cfg['lr_min'])
    optimizer = tf.keras.optimizers.SGD(learning_rate=learning_rate,
                                        momentum=cfg['momentum'])
    optimizer_arch = tf.keras.optimizers.Adam(
        learning_rate=cfg['arch_learning_rate'], beta_1=0.5, beta_2=0.999)

    # define losses function
    criterion = CrossEntropyLoss()

    # load checkpoint
    checkpoint_dir = './checkpoints/' + cfg['sub_name']
    checkpoint = tf.train.Checkpoint(step=tf.Variable(0, name='step'),
                                     optimizer=optimizer,
                                     optimizer_arch=optimizer_arch,
                                     model=sna.model,
                                     alphas_normal=sna.alphas_normal,
                                     alphas_reduce=sna.alphas_reduce,
                                     betas_normal=sna.betas_normal,
                                     betas_reduce=sna.betas_reduce)
    manager = tf.train.CheckpointManager(checkpoint=checkpoint,
                                         directory=checkpoint_dir,
                                         max_to_keep=3)
    if manager.latest_checkpoint:
        checkpoint.restore(manager.latest_checkpoint)
        print('[*] load ckpt from {} at step {}.'.format(
            manager.latest_checkpoint, checkpoint.step.numpy()))
    else:
        print("[*] training from scratch.")
    print(f"[*] searching model after {cfg['start_search_epoch']} epochs.")

    # define training step function for model
    @tf.function
    def train_step(inputs, labels):
        with tf.GradientTape() as tape:
            logits = sna.model((inputs, *sna.arch_parameters), training=True)

            losses = {}
            losses['reg'] = tf.reduce_sum(sna.model.losses)
            losses['ce'] = criterion(labels, logits)
            total_loss = tf.add_n([l for l in losses.values()])

        grads = tape.gradient(total_loss, sna.model.trainable_variables)
        grads = [(tf.clip_by_norm(grad, cfg['grad_clip'])) for grad in grads]
        optimizer.apply_gradients(zip(grads, sna.model.trainable_variables))

        return logits, total_loss, losses

    # define training step function for arch_parameters
    @tf.function
    def train_step_arch(inputs, labels):
        with tf.GradientTape() as tape:
            logits = sna.model((inputs, *sna.arch_parameters), training=True)

            losses = {}
            losses['reg'] = cfg['arch_weight_decay'] * tf.add_n(
                [tf.reduce_sum(p**2) for p in sna.arch_parameters])
            losses['ce'] = criterion(labels, logits)
            total_loss = tf.add_n([l for l in losses.values()])

        grads = tape.gradient(total_loss, sna.arch_parameters)
        optimizer_arch.apply_gradients(zip(grads, sna.arch_parameters))

        return losses

    summary_writer = tf.summary.create_file_writer('./logs/' + cfg['sub_name'])

    print("[*] finished searching for one epoch")

    print("[*] Start sampling architetures")

    prog_bar = ProgressBar(50, 0)

    # Start sampling for 50 archs
    for geno_num in range(50):
        genotype = sna.get_genotype(random_search_flag=True)
        prog_bar.update(f"\n Sampled{geno_num}th arch: {genotype}")
        # print(f"\n Sampled {geno_num}th arch: {genotype}")
        f = open(
            os.path.join('./logs', cfg['sub_name'],
                         'search_random_arch_genotype.py'), 'a')
        f.write(f"\n{cfg['sub_name']}_{geno_num} = {genotype}\n")
        f.close()

    print("Sampling done!")
    debugpy.wait_for_client()
Exemple #23
0
def main(_):
    os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
    os.environ['CUDA_VISIBLE_DEVICES'] = FLAGS.gpu

    logger = tf.get_logger()
    logger.disabled = True
    logger.setLevel(logging.FATAL)
    set_memory_growth()

    cfg = load_yaml(FLAGS.cfg_path)

    model = ArcFaceModel(size=cfg['input_size'],
                         backbone_type=cfg['backbone_type'],
                         num_classes=cfg['num_classes'],
                         head_type=cfg['head_type'],
                         embd_shape=cfg['embd_shape'],
                         w_decay=cfg['w_decay'],
                         training=True)
    model.summary(line_length=80)

    if cfg['train_dataset']:
        logging.info("load ms1m dataset.")
        dataset_len = cfg['num_samples']
        steps_per_epoch = dataset_len // cfg['batch_size']
        train_dataset = dataset.load_tfrecord_dataset(cfg['train_dataset'],
                                                      cfg['batch_size'],
                                                      cfg['binary_img'],
                                                      is_ccrop=cfg['is_ccrop'])
    else:
        logging.info("load fake dataset.")
        dataset_len = 1
        steps_per_epoch = 1
        train_dataset = dataset.load_fake_dataset(cfg['input_size'])

    learning_rate = tf.constant(cfg['base_lr'])
    optimizer = tf.keras.optimizers.SGD(learning_rate=learning_rate,
                                        momentum=0.9,
                                        nesterov=True)
    loss_fn = SoftmaxLoss()

    ckpt_path = tf.train.latest_checkpoint('./checkpoints/' + cfg['sub_name'])
    if ckpt_path is not None:
        print("[*] load ckpt from {}".format(ckpt_path))
        model.load_weights(ckpt_path)
        epochs, steps = get_ckpt_inf(ckpt_path, steps_per_epoch)
    else:
        print("[*] training from scratch.")
        epochs, steps = 1, 1

    if FLAGS.mode == 'eager_tf':
        # Eager mode is great for debugging
        # Non eager graph mode is recommended for real training
        summary_writer = tf.summary.create_file_writer('./logs/' +
                                                       cfg['sub_name'])

        train_dataset = iter(train_dataset)

        while epochs <= cfg['epochs']:
            inputs, labels = next(train_dataset)

            with tf.GradientTape() as tape:
                logist = model(inputs, training=True)
                reg_loss = tf.reduce_sum(model.losses)
                pred_loss = loss_fn(labels, logist)
                total_loss = pred_loss + reg_loss

            grads = tape.gradient(total_loss, model.trainable_variables)
            optimizer.apply_gradients(zip(grads, model.trainable_variables))

            if steps % 5 == 0:
                verb_str = "Epoch {}/{}: {}/{}, loss={:.2f}, lr={:.4f}"
                print(
                    verb_str.format(epochs, cfg['epochs'],
                                    steps % steps_per_epoch, steps_per_epoch,
                                    total_loss.numpy(), learning_rate.numpy()))

                with summary_writer.as_default():
                    tf.summary.scalar('loss/total loss',
                                      total_loss,
                                      step=steps)
                    tf.summary.scalar('loss/pred loss', pred_loss, step=steps)
                    tf.summary.scalar('loss/reg loss', reg_loss, step=steps)
                    tf.summary.scalar('learning rate',
                                      optimizer.lr,
                                      step=steps)

            if steps % cfg['save_steps'] == 0:
                print('[*] save ckpt file!')
                model.save_weights('checkpoints/{}/e_{}_b_{}.ckpt'.format(
                    cfg['sub_name'], epochs, steps % steps_per_epoch))

            steps += 1
            epochs = steps // steps_per_epoch + 1
    else:
        model.compile(optimizer=optimizer,
                      loss=loss_fn,
                      run_eagerly=(FLAGS.mode == 'eager_fit'))

        mc_callback = ModelCheckpoint(
            'checkpoints/' + cfg['sub_name'] + '/e_{epoch}_b_{batch}.ckpt',
            save_freq=cfg['save_steps'] * cfg['batch_size'],
            verbose=1,
            save_weights_only=True)
        tb_callback = TensorBoard(log_dir='logs/',
                                  update_freq=cfg['batch_size'] * 5,
                                  profile_batch=0)
        tb_callback._total_batches_seen = steps
        tb_callback._samples_seen = steps * cfg['batch_size']
        callbacks = [mc_callback, tb_callback]

        history = model.fit(train_dataset,
                            epochs=cfg['epochs'],
                            steps_per_epoch=steps_per_epoch,
                            callbacks=callbacks,
                            initial_epoch=epochs - 1)

    print("[*] training done!")
Exemple #24
0
def main(_argv):
    # init
    os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
    os.environ['CUDA_VISIBLE_DEVICES'] = FLAGS.gpu

    logger = tf.get_logger()
    logger.disabled = True
    logger.setLevel(logging.FATAL)
    set_memory_growth()

    cfg = load_yaml(FLAGS.cfg_path)

    # define network
    model = RetinaFaceModel(cfg,
                            training=False,
                            iou_th=FLAGS.iou_th,
                            score_th=FLAGS.score_th)

    # load model from weights.h5
    # model.load_weights('./model/mbv2_weights.h5', by_name=True, skip_mismatch=True)

    # load checkpoint
    checkpoint_dir = './checkpoints/' + cfg['sub_name']
    checkpoint = tf.train.Checkpoint(model=model)
    if tf.train.latest_checkpoint(checkpoint_dir):
        checkpoint.restore(tf.train.latest_checkpoint(checkpoint_dir))
        print("[*] load ckpt from {}.".format(
            tf.train.latest_checkpoint(checkpoint_dir)))
    else:
        print("[*] Cannot find ckpt from {}.".format(checkpoint_dir))
        exit()

    if not FLAGS.webcam:
        if not os.path.exists(FLAGS.img_path):
            print(f"cannot find image path from {FLAGS.img_path}")
            exit()

        print("[*] Processing on single image {}".format(FLAGS.img_path))

        img_raw = cv2.imread(FLAGS.img_path)
        img = np.float32(img_raw.copy())

        # testing scale
        target_size = 320
        img_size_max = np.max(img.shape[0:2])
        resize = float(target_size) / float(img_size_max)
        img = cv2.resize(img,
                         None,
                         None,
                         fx=resize,
                         fy=resize,
                         interpolation=cv2.INTER_LINEAR)

        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

        # pad input image to avoid unmatched shape problem
        img, pad_params = pad_input_image(img, max_steps=max(cfg['steps']))

        # run model
        outputs = model(img[np.newaxis, ...]).numpy()

        # recover padding effect
        outputs = recover_pad_output(outputs, pad_params)

        img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)

        # draw and save results
        save_img_path = os.path.join('out_' + os.path.basename(FLAGS.img_path))
        for prior_index in range(len(outputs)):
            draw_bbox_landm(img, outputs[prior_index], target_size,
                            target_size)
        cv2.imwrite(save_img_path, img)
        print(f"[*] save result at {save_img_path}")

    else:
        cam = cv2.VideoCapture('./data/lichaochao.mp4')
        # cam = cv2.VideoCapture(0)
        frame_height = int(cam.get(cv2.CAP_PROP_FRAME_HEIGHT))
        frame_width = int(cam.get(cv2.CAP_PROP_FRAME_WIDTH))

        fourcc = cv2.VideoWriter_fourcc(*'XVID')
        fps = cam.get(cv2.CAP_PROP_FPS)
        out = cv2.VideoWriter('chaochao1.mp4',
                              fourcc,
                              fps=fps,
                              frameSize=(frame_height, frame_width))

        resize = FLAGS.down_scale_factor
        frame_height *= resize
        frame_width *= resize

        max_steps = max(cfg['steps'])
        img_pad_h = max_steps - frame_height % max_steps if frame_height % max_steps > 0 else 0
        img_pad_w = max_steps - frame_width % max_steps if frame_width % max_steps > 0 else 0
        priors = prior_box_tf(
            (frame_height + img_pad_h, frame_width + img_pad_w),
            cfg['min_sizes'], cfg['steps'], cfg['clip'])

        frame_index = 0
        outputs = []
        start_time = time.time()
        while cam.isOpened():
            _, frame = cam.read()
            if frame is None:
                print('no cam')
                break
            if frame_index < 5:
                frame_index += 1
                # continue
            else:
                frame_index = 0

                img = np.float32(frame.copy())
                if resize < 1:
                    img = cv2.resize(img, (0, 0),
                                     fx=resize,
                                     fy=resize,
                                     interpolation=cv2.INTER_LINEAR)
                img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

                # pad input image to avoid unmatched shape problem
                img, pad_params = pad_input_image(img, max_steps=max_steps)

                # run model
                outputs = model(img[np.newaxis, ...])

                preds = tf.concat([
                    outputs[0][0], outputs[1][0, :, 1][..., tf.newaxis],
                    outputs[2][0, :, 1][..., tf.newaxis]
                ], -1)

                decode_preds = decode_tf(preds, priors, cfg['variances'])

                selected_indices = tf.image.non_max_suppression(
                    boxes=decode_preds[:, :4],
                    scores=decode_preds[:, -1],
                    max_output_size=tf.shape(decode_preds)[0],
                    iou_threshold=FLAGS.iou_th,
                    score_threshold=FLAGS.score_th)

                outputs = tf.gather(decode_preds, selected_indices).numpy()

                # recover padding effect
                outputs = recover_pad_output(outputs,
                                             pad_params,
                                             resize=resize)

                # calculate fps
                # fps_str = "FPS: %.2f" % (1 / (time.time() - start_time))
                # start_time = time.time()
                # cv2.putText(frame, fps_str, (25, 50),
                #             cv2.FONT_HERSHEY_DUPLEX, 0.75, (0, 0, 255), 2)

            # draw results
            for prior_index in range(len(outputs)):
                draw_bbox_landm(frame, outputs[prior_index], frame_height,
                                frame_width)

            # calculate fps
            # fps_str = "FPS: %.2f" % (1 / (time.time() - start_time))
            # start_time = time.time()
            # cv2.putText(frame, fps_str, (25, 25),
            #             cv2.FONT_HERSHEY_DUPLEX, 0.75, (0, 255, 0), 2)

            # show frame
            out.write(frame)
            cv2.imshow('frame', frame)
            if cv2.waitKey(1) == ord('q'):
                exit()
def main(_):
    # init
    os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
    os.environ['CUDA_VISIBLE_DEVICES'] = FLAGS.gpu

    logger = tf.get_logger()
    logger.disabled = True
    logger.setLevel(logging.FATAL)
    set_memory_growth()

    cfg = load_yaml(FLAGS.cfg_path)

    # define network
    model = RetinaFaceModel(cfg, training=True)
    model.summary(line_length=80)

    # define prior box
    priors = prior_box((cfg['input_size'], cfg['input_size']),
                       cfg['min_sizes'], cfg['steps'], cfg['clip'])

    # load dataset
    train_dataset = load_dataset(cfg, priors, shuffle=True)

    # define optimizer
    steps_per_epoch = cfg['dataset_len'] // cfg['batch_size']
    learning_rate = MultiStepWarmUpLR(
        initial_learning_rate=cfg['init_lr'],
        lr_steps=[e * steps_per_epoch for e in cfg['lr_decay_epoch']],
        lr_rate=cfg['lr_rate'],
        warmup_steps=cfg['warmup_epoch'] * steps_per_epoch,
        min_lr=cfg['min_lr'])
    optimizer = tf.keras.optimizers.SGD(learning_rate=learning_rate,
                                        momentum=0.9,
                                        nesterov=True)

    # define losses function
    multi_box_loss = MultiBoxLoss()

    # load checkpoint
    checkpoint_dir = '/content/drive/My Drive/Colab/checkpoints/' + cfg[
        'sub_name']
    checkpoint = tf.train.Checkpoint(step=tf.Variable(0, name='step'),
                                     optimizer=optimizer,
                                     model=model)
    manager = tf.train.CheckpointManager(checkpoint=checkpoint,
                                         directory=checkpoint_dir,
                                         max_to_keep=3)
    if manager.latest_checkpoint:
        checkpoint.restore(manager.latest_checkpoint)
        print('[*] load ckpt from {} at step {}.'.format(
            manager.latest_checkpoint, checkpoint.step.numpy()))
    else:
        print("[*] training from scratch.")

    # define training step function
    @tf.function
    def train_step(inputs, labels):
        with tf.GradientTape() as tape:
            predictions = model(inputs, training=True)

            losses = {}
            losses['reg'] = tf.reduce_sum(model.losses)
            losses['loc'], losses['landm'], losses['class'] = \
                multi_box_loss(labels, predictions)
            total_loss = tf.add_n([l for l in losses.values()])

        grads = tape.gradient(total_loss, model.trainable_variables)
        optimizer.apply_gradients(zip(grads, model.trainable_variables))

        return total_loss, losses

    # training loop
    summary_writer = tf.summary.create_file_writer('./logs/' + cfg['sub_name'])
    remain_steps = max(
        steps_per_epoch * cfg['epoch'] - checkpoint.step.numpy(), 0)
    prog_bar = ProgressBar(steps_per_epoch,
                           checkpoint.step.numpy() % steps_per_epoch)

    for inputs, labels in train_dataset.take(remain_steps):
        checkpoint.step.assign_add(1)
        steps = checkpoint.step.numpy()

        total_loss, losses = train_step(inputs, labels)

        prog_bar.update("epoch={}/{}, loss={:.4f}, lr={:.1e}".format(
            ((steps - 1) // steps_per_epoch) + 1, cfg['epoch'],
            total_loss.numpy(),
            optimizer.lr(steps).numpy()))

        if steps % 10 == 0:
            with summary_writer.as_default():
                tf.summary.scalar('loss/total_loss', total_loss, step=steps)
                for k, l in losses.items():
                    tf.summary.scalar('loss/{}'.format(k), l, step=steps)
                tf.summary.scalar('learning_rate',
                                  optimizer.lr(steps),
                                  step=steps)

        if steps % cfg['save_steps'] == 0:
            manager.save()
            print("\n[*] save ckpt file at {}".format(
                manager.latest_checkpoint))

    manager.save()
    print("\n[*] training done! save ckpt file at {}".format(
        manager.latest_checkpoint))
Exemple #26
0
def main(_argv):
    mkdir(FLAGS.destination_dir)
    os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
    os.environ['CUDA_VISIBLE_DEVICES'] = FLAGS.gpu

    logger = tf.get_logger()
    logger.disabled = True
    logger.setLevel(logging.FATAL)
    set_memory_growth()

    cfg = load_yaml(FLAGS.cfg_path)
    aligner = FaceAligner(desiredFaceSize=128)
    # define network
    model = RetinaFaceModel(cfg, training=False, iou_th=FLAGS.iou_th,
                            score_th=FLAGS.score_th)

    # load checkpoint
    checkpoint_dir = './checkpoints/' + cfg['sub_name']
    checkpoint = tf.train.Checkpoint(model=model)
    if tf.train.latest_checkpoint(checkpoint_dir):
        checkpoint.restore(tf.train.latest_checkpoint(checkpoint_dir))
        print("[*] load ckpt from {}.".format(
            tf.train.latest_checkpoint(checkpoint_dir)))
    else:
        print("[*] Cannot find ckpt from {}.".format(checkpoint_dir))
        exit()
    total = 0
    processed_total = 0
    CLASS_NAMES = np.array(os.listdir(FLAGS.folder_path))
    temp = os.listdir(FLAGS.destination_dir)
    temp.sort()
    CLASS_NAMES.sort()
    for f in CLASS_NAMES:
        processed_image = 0
        ######################################
        # Need modified for using
        ######################################
        if os.path.isfile(FLAGS.folder_path+f):
            continue
        if f in temp and f != temp[-1]:
          continue
        items = os.listdir(FLAGS.folder_path+f)
        mkdir(FLAGS.destination_dir+f)
        for path in items:
            frame = cv2.imread(FLAGS.folder_path + f +'/'+ path)
            if frame is None:
              continue
            frame_height, frame_width, _ = frame.shape
            img = np.float32(frame.copy())
            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

            # pad input image to avoid unmatched shape problem
            img, pad_params = pad_input_image(img, max_steps=max(cfg['steps']))

            # run model
            outputs = model(img[np.newaxis, ...]).numpy()

            # recover padding effect
            outputs = recover_pad_output(outputs, pad_params)
            if len(outputs) < 1:
              continue
            ann = max(outputs, key=lambda x: (x[2]-x[0])*(x[3]-x[1]))
            b_box = int(ann[0] * frame_width), int(ann[1] * frame_height), \
                    int(ann[2] * frame_width), int(ann[3] * frame_height)
            if (b_box[0]<0) or (b_box[1]<0) or (b_box[2]>=frame_width) or (b_box[3]>=frame_height):
              continue
            keypoints = {
                'left_eye': (ann[4] * frame_width,ann[5] * frame_height),
                'right_eye': (ann[6] * frame_width,ann[7] * frame_height),
                'nose': (ann[8], ann[9]),
                'left_mouth': (ann[10] * frame_width, ann[11] * frame_height),
                'right_mouth': (ann[12] * frame_width,ann[13] * frame_height),
            }
            # croped_image = frame[b_box[1]:b_box[3],b_box[0]:b_box[2], :]
            # out_frame = cv2.resize(croped_image, (112,112), interpolation=cv2.INTER_CUBIC)
            out_frame = aligner.align(frame, keypoints, b_box)
            # for i in range(4,14):
            #     if i%2 == 0:
            #         ann[i] = int(ann[i]*frame_width)
            #     else:
            #         ann[i] = int(ann[i]*frame_height)
            # out_frame = norm_crop(frame, np.array([ann[4:6],ann[6:8],ann[8:10],ann[10:12],ann[12:14]]))
            try:
                cv2.imwrite(FLAGS.destination_dir + f +'/'+ path, out_frame)
                processed_image += 1
            except FileExistsError as e:
                pass
        
        print(f + " Done")
        total += len(items)
        processed_total += processed_image
Exemple #27
0
def main(_):
    # init
    os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
    os.environ['CUDA_VISIBLE_DEVICES'] = FLAGS.gpu
    set_memory_growth()

    # load PRNet model
    cfg = load_yaml(FLAGS.cfg_path)
    model = PRN(cfg, is_dlib=True)

    # evaluation
    if not FLAGS.use_cam:  # on test-img
        print(
            "[*] Processing on images in {}. Press 's' to save result.".format(
                FLAGS.img_path))
        img_paths = glob.glob(os.path.join(FLAGS.img_path, '*'))
        for img_path in img_paths:
            img = cv2.imread(img_path)
            pos = model.process(img_path)
            if pos is None:
                continue

            vertices = model.get_vertices(pos)
            kpt = model.get_landmarks(pos)
            camera_matrix, _ = estimate_pose(vertices)

            cv2.imshow('Input', img)
            cv2.imshow('Sparse alignment', plot_kpt(img, kpt))
            cv2.imshow('Dense alignment', plot_vertices(img, vertices))
            cv2.imshow('Pose', plot_pose_box(img, camera_matrix, kpt))
            cv2.moveWindow('Input', 0, 0)
            cv2.moveWindow('Sparse alignment', 500, 0)
            cv2.moveWindow('Dense alignment', 1000, 0)
            cv2.moveWindow('Pose', 1500, 0)

            key = cv2.waitKey(0)
            if key == ord('q'):
                exit()
            elif key == ord('s'):
                cv2.imwrite(
                    os.path.join(FLAGS.save_path, os.path.basename(img_path)),
                    plot_kpt(img, kpt))
                print("Result saved in {}".format(FLAGS.save_path))

    else:  # webcam demo
        cap = cv2.VideoCapture(0)
        start_time = time.time()
        count = 1
        while (True):
            _, image = cap.read()

            pos = model.process(image)
            fps_str = 'FPS: %.2f' % (1 / (time.time() - start_time))
            start_time = time.time()
            cv2.putText(image, fps_str, (25, 25), cv2.FONT_HERSHEY_DUPLEX,
                        0.75, (0, 255, 0), 2)
            cv2.imshow('Input', image)
            cv2.moveWindow('Input', 0, 0)

            key = cv2.waitKey(1)
            if pos is None:
                cv2.waitKey(1)
                cv2.destroyWindow('Sparse alignment')
                cv2.destroyWindow('Dense alignment')
                cv2.destroyWindow('Pose')
                if key & 0xFF == ord('q'):
                    break
                continue

            else:
                vertices = model.get_vertices(pos)
                kpt = model.get_landmarks(pos)
                camera_matrix, _ = estimate_pose(vertices)

                result_list = [
                    plot_kpt(image, kpt),
                    plot_vertices(image, vertices),
                    plot_pose_box(image, camera_matrix, kpt)
                ]

                cv2.imshow('Sparse alignment', result_list[0])
                cv2.imshow('Dense alignment', result_list[1])
                cv2.imshow('Pose', result_list[2])
                cv2.moveWindow('Sparse alignment', 500, 0)
                cv2.moveWindow('Dense alignment', 1000, 0)
                cv2.moveWindow('Pose', 1500, 0)

                if key & 0xFF == ord('s'):
                    image_name = 'prnet_cam_' + str(count)
                    save_path = FLAGS.save_path

                    cv2.imwrite(
                        os.path.join(save_path, image_name + '_result.jpg'),
                        np.concatenate(result_list, axis=1))
                    cv2.imwrite(
                        os.path.join(save_path, image_name + '_image.jpg'),
                        image)
                    count += 1
                    print("Result saved in {}".format(FLAGS.save_path))

                if key & 0xFF == ord('q'):
                    break
Exemple #28
0
def main(_):
    # init
    os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
    os.environ['CUDA_VISIBLE_DEVICES'] = FLAGS.gpu

    logger = tf.get_logger()
    logger.disabled = True
    logger.setLevel(logging.FATAL)
    set_memory_growth()

    cfg = load_yaml(FLAGS.cfg_path)

    # define training step function
    @tf.function
    def train_step(inputs, labels, drop_path_prob):
        with tf.GradientTape() as tape:
            logits, logits_aux = model((inputs, drop_path_prob), training=True)

            losses = {}
            losses['reg'] = tf.reduce_sum(model.losses)
            losses['ce'] = criterion(labels, logits)
            losses['ce_auxiliary'] = \
                cfg['auxiliary_weight'] * criterion(labels, logits_aux)
            total_loss = tf.add_n([l for l in losses.values()])

        grads = tape.gradient(total_loss, model.trainable_variables)
        grads = [(tf.clip_by_norm(grad, cfg['grad_clip'])) for grad in grads]
        optimizer.apply_gradients(zip(grads, model.trainable_variables))

        return logits, total_loss, losses

    # Used to store the final accuracy for every arch
    final_acc = pd.DataFrame(data=None, columns=['arch_name', 'acc'])

    loop_num = 50

    if Debug:
        # debugpy.wait_for_client()
        loop_num = 1
    # define network
    for arch_num in range(loop_num):
        # read the arch
        arch = str(f"{cfg['sub_name']}_{arch_num}")
        cfg['arch'] = arch

        model = CifarModel(cfg, training=True, file_name=FLAGS.file_name)
        if Debug:
            model.summary(line_length=80)
            print("param size = {:f}MB".format(count_parameters_in_MB(model)))

        # load dataset
        train_dataset = load_cifar10_dataset(
            cfg['batch_size'],
            split='train',
            shuffle=True,
            drop_remainder=True,
            using_normalize=cfg['using_normalize'],
            using_crop=cfg['using_crop'],
            using_flip=cfg['using_flip'],
            using_cutout=cfg['using_cutout'],
            cutout_length=cfg['cutout_length'])
        val_dataset = load_cifar10_dataset(
            cfg['val_batch_size'],
            split='test',
            shuffle=False,
            drop_remainder=False,
            using_normalize=cfg['using_normalize'],
            using_crop=False,
            using_flip=False,
            using_cutout=False)

        # define optimizer
        steps_per_epoch = cfg['dataset_len'] // cfg['batch_size']
        learning_rate = CosineAnnealingLR(initial_learning_rate=cfg['init_lr'],
                                          t_period=cfg['epoch'] *
                                          steps_per_epoch,
                                          lr_min=cfg['lr_min'])
        optimizer = tf.keras.optimizers.SGD(learning_rate=learning_rate,
                                            momentum=cfg['momentum'])

        # define losses function
        criterion = CrossEntropyLoss()

        # load checkpoint
        checkpoint_dir = './checkpoints/' + arch
        checkpoint = tf.train.Checkpoint(step=tf.Variable(0, name='step'),
                                         optimizer=optimizer,
                                         model=model)
        manager = tf.train.CheckpointManager(checkpoint=checkpoint,
                                             directory=checkpoint_dir,
                                             max_to_keep=3)
        if manager.latest_checkpoint:
            checkpoint.restore(manager.latest_checkpoint)
            print('[*] load ckpt from {} at step {}.'.format(
                manager.latest_checkpoint, checkpoint.step.numpy()))
        else:
            print("[*] training from scratch.")

        # training loop
        summary_writer = tf.summary.create_file_writer('./logs/' +
                                                       cfg['sub_name'])
        total_steps = steps_per_epoch * cfg['epoch']
        remain_steps = max(total_steps - checkpoint.step.numpy(), 0)
        prog_bar = ProgressBar(steps_per_epoch,
                               checkpoint.step.numpy() % steps_per_epoch)

        train_acc = AvgrageMeter()
        val_acc = AvgrageMeter()
        best_acc = 0.
        for inputs, labels in train_dataset.take(remain_steps):
            checkpoint.step.assign_add(1)
            drop_path_prob = cfg['drop_path_prob'] * (
                tf.cast(checkpoint.step, tf.float32) / total_steps)
            steps = checkpoint.step.numpy()
            epochs = ((steps - 1) // steps_per_epoch) + 1

            logits, total_loss, losses = train_step(inputs, labels,
                                                    drop_path_prob)
            train_acc.update(
                accuracy(logits.numpy(), labels.numpy())[0], cfg['batch_size'])

            prog_bar.update(
                "epoch={}/{}, loss={:.4f}, acc={:.2f}, lr={:.2e}".format(
                    epochs, cfg['epoch'], total_loss.numpy(), train_acc.avg,
                    optimizer.lr(steps).numpy()))

            if steps % cfg['val_steps'] == 0 and steps > 1:
                print("\n[*] validate...", end='')
                val_acc.reset()
                for inputs_val, labels_val in val_dataset:
                    logits_val, _ = model((inputs_val, tf.constant([0.])))
                    val_acc.update(
                        accuracy(logits_val.numpy(), labels_val.numpy())[0],
                        inputs_val.shape[0])

                if val_acc.avg > best_acc:
                    best_acc = val_acc.avg
                    model.save_weights(
                        f"checkpoints/{cfg['sub_name']}/best.ckpt")

                val_str = " val acc {:.2f}%, best acc {:.2f}%"
                print(val_str.format(val_acc.avg, best_acc), end='')

            if steps % 10 == 0:
                with summary_writer.as_default():
                    tf.summary.scalar('acc/train', train_acc.avg, step=steps)
                    tf.summary.scalar('acc/val', val_acc.avg, step=steps)

                    tf.summary.scalar('loss/total_loss',
                                      total_loss,
                                      step=steps)
                    for k, l in losses.items():
                        tf.summary.scalar('loss/{}'.format(k), l, step=steps)
                    tf.summary.scalar('learning_rate',
                                      optimizer.lr(steps),
                                      step=steps)

            if steps % cfg['save_steps'] == 0:
                manager.save()
                print("\n[*] save ckpt file at {}".format(
                    manager.latest_checkpoint))

            if steps % steps_per_epoch == 0:
                train_acc.reset()

        manager.save()
        print("\n[*] training one arch done! save ckpt file at {}".format(
            manager.latest_checkpoint))
        final_acc.loc[arch_num] = list([arch, best_acc])
    print("Whole training ended, the best result is :")
    print("\t", final_acc.iloc[final_acc['acc'].idxmax()])
Exemple #29
0
def main(_argv):
    # init
    os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"
    os.environ["CUDA_VISIBLE_DEVICES"] = FLAGS.gpu

    logger = tf.get_logger()
    logger.disabled = True
    logger.setLevel(logging.FATAL)
    set_memory_growth()

    cfg = load_yaml(FLAGS.cfg_path)

    # define network
    model = RetinaFaceModel(cfg,
                            training=False,
                            iou_th=FLAGS.iou_th,
                            score_th=FLAGS.score_th)

    # load checkpoint
    checkpoint_dir = "./checkpoints/" + cfg["sub_name"]
    checkpoint = tf.train.Checkpoint(model=model)
    if tf.train.latest_checkpoint(checkpoint_dir):
        checkpoint.restore(tf.train.latest_checkpoint(checkpoint_dir))
        print("[*] load ckpt from {}.".format(
            tf.train.latest_checkpoint(checkpoint_dir)))
    else:
        print("[*] Cannot find ckpt from {}.".format(checkpoint_dir))
        exit()

    # evaluation on testing dataset
    testset_folder = cfg["testing_dataset_path"]
    testset_list = os.path.join(testset_folder, "label.txt")

    img_paths, _ = load_info(testset_list)
    for img_index, img_path in enumerate(img_paths):
        print(" [{} / {}] det {}".format(img_index + 1, len(img_paths),
                                         img_path))
        img_raw = cv2.imread(img_path, cv2.IMREAD_COLOR)
        img_height_raw, img_width_raw, _ = img_raw.shape
        img = np.float32(img_raw.copy())

        # testing scale
        target_size = 1600
        max_size = 2150
        img_shape = img.shape
        img_size_min = np.min(img_shape[0:2])
        img_size_max = np.max(img_shape[0:2])
        resize = float(target_size) / float(img_size_min)
        # prevent bigger axis from being more than max_size:
        if np.round(resize * img_size_max) > max_size:
            resize = float(max_size) / float(img_size_max)
        if FLAGS.origin_size:
            if os.path.basename(img_path) == "6_Funeral_Funeral_6_618.jpg":
                resize = 0.5  # this image is too big to avoid OOM problem
            else:
                resize = 1

        img = cv2.resize(img,
                         None,
                         None,
                         fx=resize,
                         fy=resize,
                         interpolation=cv2.INTER_LINEAR)
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

        # pad input image to avoid unmatched shape problem
        img, pad_params = pad_input_image_(img, max_steps=max(cfg["steps"]))

        # run model
        outputs = model(img[np.newaxis, ...]).numpy()

        # recover padding effect
        outputs = recover_pad_output(outputs, pad_params)

        # write results
        img_name = os.path.basename(img_path)
        sub_dir = os.path.basename(os.path.dirname(img_path))
        save_name = os.path.join(FLAGS.save_folder, sub_dir,
                                 img_name.replace(".jpg", ".txt"))

        pathlib.Path(os.path.join(FLAGS.save_folder,
                                  sub_dir)).mkdir(parents=True, exist_ok=True)

        with open(save_name, "w") as file:
            bboxs = outputs[:, :4]
            confs = outputs[:, -1]

            file_name = img_name + "\n"
            bboxs_num = str(len(bboxs)) + "\n"
            file.write(file_name)
            file.write(bboxs_num)
            for box, conf in zip(bboxs, confs):
                x = int(box[0] * img_width_raw)
                y = int(box[1] * img_height_raw)
                w = int(box[2] * img_width_raw) - int(box[0] * img_width_raw)
                h = int(box[3] * img_height_raw) - int(box[1] * img_height_raw)
                confidence = str(conf)
                line = str(x) + " " + str(y) + " " + str(w) + " " + str(
                    h) + " " + confidence + " \n"
                file.write(line)

        # save images
        pathlib.Path(os.path.join("./results", cfg["sub_name"],
                                  sub_dir)).mkdir(parents=True, exist_ok=True)
        if FLAGS.save_image:
            for prior_index in range(len(outputs)):
                if outputs[prior_index][15] >= FLAGS.vis_th:
                    draw_bbox_landm(img_raw, outputs[prior_index],
                                    img_height_raw, img_width_raw)
            cv2.imwrite(
                os.path.join("./results", cfg["sub_name"], sub_dir, img_name),
                img_raw)
Exemple #30
0
def main(_argv):
    os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
    os.environ['CUDA_VISIBLE_DEVICES'] = FLAGS.gpu

    logger = tf.get_logger()
    logger.disabled = True
    logger.setLevel(logging.FATAL)
    set_memory_growth()

    cfg = load_yaml(FLAGS.cfg_path)

    model = ArcFaceModel(size=cfg['input_size'],
                         backbone_type=cfg['backbone_type'],
                         training=False)

    ckpt_path = tf.train.latest_checkpoint('./checkpoints/' + cfg['sub_name'])
    if ckpt_path is not None:
        print("[*] load ckpt from {}".format(ckpt_path))
        model.load_weights(ckpt_path)
    else:
        print("[*] Cannot find ckpt from {}.".format(ckpt_path))
        exit()

    if FLAGS.img_path:

        print("[*] Encode {} to ./output_embeds.npy".format(FLAGS.img_path))

        # file_dir = "C:/Users/chaehyun/PycharmProjects/ArcFace37_TF2x/data/AFDB_masked_face_dataset/*"
        file_dir = "C:/Users/smgg/Desktop/dataset/superjunior/*.jpg"
        # detector = MTCNN()
        i = 0
        img_list = glob.glob(file_dir)
        for i_list in img_list:
            img = cv2.cvtColor(cv2.imread(i_list), cv2.COLOR_BGR2RGB)
            detector = MTCNN()
            data_list = detector.detect_faces(img)

            for data in data_list:
                xmin, ymin, width, height = data['box']
                xmax = xmin + width
                ymax = ymin + height

                face_image = img[ymin:ymax, xmin:xmax, :]
                face_image = cv2.cvtColor(face_image, cv2.COLOR_RGB2BGR)
                cv2.imshow('', face_image)

                print(i_list)

                img_resize = cv2.resize(face_image,
                                        (cfg['input_size'], cfg['input_size']))
                cv2.imwrite("./{}.jpg".format(i), img_resize)
                cv2.imshow('', img_resize)
                cv2.waitKey(0)
                i = i + 1

                cv2.putText(img, "0", (xmin, ymin - 10),
                            cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 1,
                            cv2.LINE_AA)
                cv2.rectangle(img, (xmin, ymin), (xmax, ymax), (0, 255, 0), 1)

                img_resize = img_resize.astype(np.float32) / 255.
                if len(img_resize.shape) == 3:
                    img_resize = np.expand_dims(img_resize, 0)
                embeds = l2_norm(model(img_resize))

                #print('./newTWICE_id/{}.npy'.format(i_list.split('/')[6].split('\\')[1].split('.jpg')[0]))
                np.save(
                    './newTWICE_id/{}.npy'.format(
                        i_list.split('/')[5].split('\\')[1].split('.jpg')[0]),
                    embeds)
    else:
        print("[*] Loading LFW, AgeDB30 and CFP-FP...")
        lfw, agedb_30, cfp_fp, lfw_issame, agedb_30_issame, cfp_fp_issame = \
            get_val_data(cfg['test_dataset'])

        print("[*] Perform Evaluation on LFW...")
        acc_lfw, best_th = perform_val(cfg['embd_shape'],
                                       cfg['batch_size'],
                                       model,
                                       lfw,
                                       lfw_issame,
                                       is_ccrop=cfg['is_ccrop'])
        print("    acc {:.4f}, th: {:.2f}".format(acc_lfw, best_th))

        print("[*] Perform Evaluation on AgeDB30...")
        acc_agedb30, best_th = perform_val(cfg['embd_shape'],
                                           cfg['batch_size'],
                                           model,
                                           agedb_30,
                                           agedb_30_issame,
                                           is_ccrop=cfg['is_ccrop'])
        print("    acc {:.4f}, th: {:.2f}".format(acc_agedb30, best_th))

        print("[*] Perform Evaluation on CFP-FP...")
        acc_cfp_fp, best_th = perform_val(cfg['embd_shape'],
                                          cfg['batch_size'],
                                          model,
                                          cfp_fp,
                                          cfp_fp_issame,
                                          is_ccrop=cfg['is_ccrop'])
        print("    acc {:.4f}, th: {:.2f}".format(acc_cfp_fp, best_th))