chenshijun
2019-06-05 f111cb002b9c6065fdf6bb274ce5857a9e875e8c
code/test_on_video_dlib.py
@@ -58,12 +58,12 @@
    # Dlib face detection model
    cnn_face_detector = dlib.cnn_face_detection_model_v1(args.face_model)
    print 'Loading snapshot.'
    print('Loading snapshot.')
    # Load snapshot
    saved_state_dict = torch.load(snapshot_path)
    model.load_state_dict(saved_state_dict)
    print 'Loading data.'
    print ('Loading data.')
    transformations = transforms.Compose([transforms.Scale(224),
    transforms.CenterCrop(224), transforms.ToTensor(),
@@ -71,13 +71,13 @@
    model.cuda(gpu)
    print 'Ready to test network.'
    print('Ready to test network.')
    # Test the Model
    model.eval()  # Change model to 'eval' mode (BN uses moving mean/var).
    total = 0
    idx_tensor = [idx for idx in xrange(66)]
    idx_tensor = [idx for idx in range(66)]
    idx_tensor = torch.FloatTensor(idx_tensor).cuda(gpu)
    video = cv2.VideoCapture(video_path)
@@ -103,35 +103,40 @@
    frame_num = 1
    while frame_num <= args.n_frames:
        print(frame_num)
        ret,frame = video.read()
        if ret == False:
            break
        frame = cv2.resize(frame, (960, 540), interpolation=cv2.INTER_LINEAR)
        print(frame.shape)
        cv2_frame = cv2.cvtColor(frame,cv2.COLOR_BGR2RGB)
        # Dlib detect
        dets = cnn_face_detector(cv2_frame, 1)
        for idx, det in enumerate(dets):
            print('**********************yes****************************yes********************************yes************************************')
            # Get x_min, y_min, x_max, y_max, conf
            x_min = det.rect.left()
            y_min = det.rect.top()
            x_max = det.rect.right()
            y_max = det.rect.bottom()
            conf = det.confidence
            print x_min, y_min, x_max, y_max, conf
            if conf > 0.95:
            if conf > 1.0:
                bbox_width = abs(x_max - x_min)
                bbox_height = abs(y_max - y_min)
                x_min -= 3 * bbox_width / 4
                x_max += 3 * bbox_width / 4
                x_min -= 2 * bbox_width / 4
                x_max += 2 * bbox_width / 4
                y_min -= 3 * bbox_height / 4
                y_max += bbox_height / 4
                x_min = max(x_min, 0); y_min = max(y_min, 0)
                x_max = min(frame.shape[1], x_max); y_max = min(frame.shape[0], y_max)
                # Crop image
                img = cv2_frame[y_min:y_max,x_min:x_max]
                img = cv2_frame[int(y_min):int(y_max),int(x_min):int(x_max)]
                img = Image.fromarray(img)
                # Transform
@@ -157,8 +162,8 @@
                # Plot expanded bounding box
                # cv2.rectangle(frame, (x_min, y_min), (x_max, y_max), (0,255,0), 1)
            out.write(frame)
            frame_num += 1
        out.write(frame)
        frame_num += 1
    out.release()
    video.release()