From 3e0d6320a4acdb3c73e0efb31ee57fb1534237b2 Mon Sep 17 00:00:00 2001
From: chenshijun <csj_sky@126.com>
Date: 星期三, 05 六月 2019 10:41:06 +0800
Subject: [PATCH] real time video
---
code/test_on_video_dlib.py | 6 +
demo.txt | 3
code/demo2.py | 168 ++++++++++++++++++++++++++++++++++++++++++
code/utils.py | 2
test/demo.py | 31 +++++++
test/demo4.py | 0
test/demo3.py | 15 +++
7 files changed, 222 insertions(+), 3 deletions(-)
diff --git a/code/demo2.py b/code/demo2.py
new file mode 100644
index 0000000..be61416
--- /dev/null
+++ b/code/demo2.py
@@ -0,0 +1,168 @@
+import sys, os, argparse
+
+import numpy as np
+import cv2
+import matplotlib.pyplot as plt
+
+import torch
+import torch.nn as nn
+from torch.autograd import Variable
+from torch.utils.data import DataLoader
+from torchvision import transforms
+import torch.backends.cudnn as cudnn
+import torchvision
+import torch.nn.functional as F
+from PIL import Image
+
+import datasets, hopenet, utils
+
+from skimage import io
+import dlib
+import time
+
+if __name__ == '__main__':
+
+
+ cudnn.enabled = True
+
+ batch_size = 1
+ gpu = 'cuda:0'
+ snapshot_path = './snapshots/hopenet_robust_alpha1.pkl'
+ out_dir = 'output/video'
+
+
+ if not os.path.exists(out_dir):
+ os.makedirs(out_dir)
+
+ # ResNet50 structure
+ model = hopenet.Hopenet(torchvision.models.resnet.Bottleneck, [3, 4, 6, 3], 66)
+
+ # Dlib face detection model
+ cnn_face_detector = dlib.cnn_face_detection_model_v1('./models/mmod_human_face_detector.dat')
+
+ print('Loading snapshot.')
+ # Load snapshot
+ saved_state_dict = torch.load(snapshot_path)
+ model.load_state_dict(saved_state_dict)
+
+ print ('Loading data.')
+
+ transformations = transforms.Compose([transforms.Scale(224),
+ transforms.CenterCrop(224), transforms.ToTensor(),
+ transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])])
+
+ model.cuda(gpu)
+
+ 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 range(66)]
+ idx_tensor = torch.FloatTensor(idx_tensor).cuda(gpu)
+
+ video = cv2.VideoCapture('rtsp://admin:a1234567@192.168.1.215:554/H.264/ch1/main/av_stream')
+ if not video.isOpened():
+ print('not opened')
+ exit()
+
+ # New cv2
+ width = int(video.get(cv2.CAP_PROP_FRAME_WIDTH)) # float
+ height = int(video.get(cv2.CAP_PROP_FRAME_HEIGHT)) # float
+ print('width=========='+str(width))
+ print('height=========='+str(height))
+ # Define the codec and create VideoWriter object
+ fourcc = cv2.VideoWriter_fourcc(*'MJPG')
+ # out = cv2.VideoWriter('output/video/output-%s.avi' % 's', fourcc, 25, (width, height))
+ out = cv2.VideoWriter('output/video/output-%s.avi' % 's', fourcc, 25, (1024, 534))
+ # # Old cv2
+ # width = int(video.get(cv2.cv.CV_CAP_PROP_FRAME_WIDTH)) # float
+ # height = int(video.get(cv2.cv.CV_CAP_PROP_FRAME_HEIGHT)) # float
+ #
+ # # Define the codec and create VideoWriter object
+ # fourcc = cv2.cv.CV_FOURCC(*'MJPG')
+ # out = cv2.VideoWriter('output/video/output-%s.avi' % args.output_string, fourcc, 30.0, (width, height))
+
+ txt_out = open('/home/basic/code/deep-head-pose/output/images/%s.txt' % 's', 'w')
+
+ frame_num = 1
+ start = time.time()
+ while frame_num <= 10000:
+ # print(frame_num)
+
+
+ ret,frame = video.read()
+ if ret == False:
+ break
+
+ frame = cv2.resize(frame, (1024, 534), interpolation=cv2.INTER_LINEAR)
+ if frame_num % 8 != 0:
+ frame_num += 1
+ cv2.imshow("frame", frame)
+ if cv2.waitKey(1) & 0xFF == ord('q'):
+ break
+ # out.write(frame)
+ continue
+ 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
+
+ if conf > 1.0:
+ bbox_width = abs(x_max - x_min)
+ bbox_height = abs(y_max - y_min)
+ 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[int(y_min):int(y_max),int(x_min):int(x_max)]
+ img = Image.fromarray(img)
+
+ # Transform
+ img = transformations(img)
+ img_shape = img.size()
+ img = img.view(1, img_shape[0], img_shape[1], img_shape[2])
+ img = Variable(img).cuda(gpu)
+
+ yaw, pitch, roll = model(img)
+
+ yaw_predicted = F.softmax(yaw)
+ pitch_predicted = F.softmax(pitch)
+ roll_predicted = F.softmax(roll)
+ # Get continuous predictions in degrees.
+ yaw_predicted = torch.sum(yaw_predicted.data[0] * idx_tensor) * 3 - 99
+ pitch_predicted = torch.sum(pitch_predicted.data[0] * idx_tensor) * 3 - 99
+ roll_predicted = torch.sum(roll_predicted.data[0] * idx_tensor) * 3 - 99
+
+ # Print new frame with cube and axis
+ txt_out.write(str(frame_num) + ' %f %f %f\n' % (yaw_predicted, pitch_predicted, roll_predicted))
+ # utils.plot_pose_cube(frame, yaw_predicted, pitch_predicted, roll_predicted, (x_min + x_max) / 2, (y_min + y_max) / 2, size = bbox_width)
+ # utils.draw_axis(frame, yaw_predicted, pitch_predicted, roll_predicted, tdx = (x_min + x_max) / 2, tdy= (y_min + y_max) / 2, size = bbox_height/2)
+ utils.draw_axis(frame, yaw_predicted, pitch_predicted, roll_predicted, tdx = (x_min + x_max)/2, tdy= (y_min + y_max)/2, size = bbox_height*3)
+ # Plot expanded bounding box
+ # cv2.rectangle(frame, (x_min, y_min), (x_max, y_max), (0,255,0), 1)
+ cv2.imshow("frame", frame)
+ if cv2.waitKey(1) & 0xFF == ord('q'):
+ break
+ # out.write(frame)
+ frame_num += 1
+ end = time.time()
+
+ print(end - start)
+ out.release()
+ video.release()
+ # cv2.destroyAllWindows()
diff --git a/code/test_on_video_dlib.py b/code/test_on_video_dlib.py
index eefced7..8ae74fa 100644
--- a/code/test_on_video_dlib.py
+++ b/code/test_on_video_dlib.py
@@ -106,15 +106,19 @@
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()
diff --git a/code/utils.py b/code/utils.py
index b75d51f..ff98fd3 100644
--- a/code/utils.py
+++ b/code/utils.py
@@ -112,6 +112,6 @@
cv2.line(img, (int(tdx), int(tdy)), (int(x1),int(y1)),(0,0,255),3)
cv2.line(img, (int(tdx), int(tdy)), (int(x2),int(y2)),(0,255,0),3)
- cv2.line(img, (int(tdx), int(tdy)), (int(x3),int(y3)),(255,0,0),2)
+ cv2.line(img, (int(tdx), int(tdy)), (int(x3),int(y3)),(0,0,0),5)
return img
diff --git a/demo.txt b/demo.txt
index 8aa7957..3d6ba0a 100644
--- a/demo.txt
+++ b/demo.txt
@@ -1 +1,2 @@
-python code/test_on_video_dlib.py --snapshot ./snapshots/hopenet_robust_alpha1.pkl --face_model ./models/mmod_human_face_detector.dat --video ./videos/out2.mp4 --output_string "something you want to add to the output video name and result file name" --n_frames 500 --fps 30
\ No newline at end of file
+python code/test_on_video_dlib.py --snapshot ./snapshots/hopenet_robust_alpha1.pkl --face_model ./models/mmod_human_face_detector.dat --video ./videos/out2.mp4 --output_string "something you want to add to the output video name and result file name" --n_frames 500 --fps 30
+python code/test_on_video_dlib.py --snapshot ./snapshots/hopenet_robust_alpha1.pkl --face_model ./models/mmod_human_face_detector.dat --video ./output/test/output__sssss.avi --output_string "something you want to add to the output video name and result file name" --n_frames 500 --fps 30
\ No newline at end of file
diff --git a/test/demo.py b/test/demo.py
new file mode 100644
index 0000000..a823418
--- /dev/null
+++ b/test/demo.py
@@ -0,0 +1,31 @@
+import cv2
+
+cap = cv2.VideoCapture('rtsp://admin:a1234567@192.168.1.188:554/H.264/ch1/main/av_stream')
+if not cap.isOpened():
+ print('not opened')
+ exit()
+
+# New cv2
+width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) # float
+height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) # float
+
+# Define the codec and create VideoWriter object
+fourcc = cv2.VideoWriter_fourcc(*'MJPG')
+out = cv2.VideoWriter('output/test/output-%s.avi' % 's', fourcc, 30, (width, height))
+
+i = 0
+while True:
+ ret, frame = cap.read()
+ i += 1
+ if ret:
+ print(i)
+ # frame = cv2.flip(frame, 0)
+ # cv2.imwrite('/home/basic/code/deep-head-pose/test/img_%s.png'%str(i), frame)
+ if i > 100:
+ # frame = cv2.flip(frame, 0)
+ break
+ out.write(frame)
+
+out.release()
+cap.release()
+# cv2.destroyAllWindows()
\ No newline at end of file
diff --git a/test/demo3.py b/test/demo3.py
new file mode 100644
index 0000000..d79095e
--- /dev/null
+++ b/test/demo3.py
@@ -0,0 +1,15 @@
+import multiprocessing
+
+def readFrame(frames):
+ frame = 0
+ frames.append(frame)
+ while frame < 1000:
+ frame += 1
+ frames.append(frame)
+
+def processFrame(frames):
+ pass
+
+frames = []
+p1 = multiprocessing.Process(target=readFrame, args=frames)
+p2 = multiprocessing.Process(target=processFrame, args=frames)
diff --git a/test/demo4.py b/test/demo4.py
new file mode 100644
index 0000000..e69de29
--- /dev/null
+++ b/test/demo4.py
--
Gitblit v1.8.0