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