Add Complete Support for Live Inferencing

lib/dataset/live_inference.py

@@ -0,0 +1,71 @@
+# -*- coding: utf-8 -*-
+
+# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is
+# holder of all proprietary rights on this computer program.
+# You can only use this computer program if you have closed
+# a license agreement with MPG or you get the right to use the computer
+# program from someone who is authorized to grant you that right.
+# Any use of the computer program without a valid license is prohibited and
+# liable to prosecution.
+#
+# Copyright©2019 Max-Planck-Gesellschaft zur Förderung
+# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute
+# for Intelligent Systems. All rights reserved.
+#
+# Contact: [email protected]
+
+import os
+import cv2
+import numpy as np
+import os.path as osp
+from torch.utils.data import Dataset
+from torchvision.transforms.functional import to_tensor
+
+from lib.utils.smooth_bbox import get_all_bbox_params
+from lib.data_utils.img_utils import get_single_image_crop_demo
+
+import time
+
+
+class LiveInference(Dataset):
+    def __init__(self, images, frames, bboxes=None, joints2d=None, scale=1.0, crop_size=224):
+
+        self.images = images
+        self.bboxes = bboxes
+        self.joints2d = joints2d
+        self.scale = scale
+        self.crop_size = crop_size
+        self.frames = frames
+        self.has_keypoints = True if joints2d is not None else False
+
+        self.norm_joints2d = np.zeros_like(self.joints2d)
+
+        if self.has_keypoints:
+            bboxes, time_pt1, time_pt2 = get_all_bbox_params(joints2d, vis_thresh=0.3)
+            bboxes[:, 2:] = 150. / bboxes[:, 2:]
+            self.bboxes = np.stack([bboxes[:, 0], bboxes[:, 1], bboxes[:, 2], bboxes[:, 2]]).T
+
+            self.images = self.images[time_pt1:time_pt2]
+            self.joints2d = joints2d[time_pt1:time_pt2]
+            self.frames = frames[time_pt1:time_pt2]
+
+
+    def __len__(self):
+        return len(self.images)
+
+    def __getitem__(self, idx):
+
+        bbox = self.bboxes[idx]
+
+        j2d = self.joints2d[idx] if self.has_keypoints else None
+
+        norm_img, raw_img, kp_2d = get_single_image_crop_demo(
+            self.images[idx],
+            bbox,
+            kp_2d=j2d,
+            scale=self.scale,
+            crop_size=self.crop_size)
+        if self.has_keypoints:
+            return norm_img, kp_2d
+        else:
+            return norm_img

lib/models/vibe.py

@@ -128,12 +128,15 @@ def __init__(
             bidirectional=False,
             use_residual=True,
             pretrained=osp.join(VIBE_DATA_DIR, 'spin_model_checkpoint.pth.tar'),
+            live_inference = False
     ):
 
         super(VIBE_Demo, self).__init__()
 
         self.seqlen = seqlen
         self.batch_size = batch_size
+        self.live_inference = live_inference
+        self.previous_feature = None
 
         self.encoder = TemporalEncoder(
             n_layers=n_layers,
@@ -157,11 +160,28 @@ def __init__(
             print(f'=> loaded pretrained model from \'{pretrained}\'')
 
 
+    # ========= Use previously stored feature vectors if live_inferencing is turned on ========= #
+    def generate_complete_feature_vector(self,feature):
+        if self.previous_feature is not None:
+            feature = torch.cat((self.previous_feature,feature),axis = 0)
+            self.previous_feature = feature[-(self.seqlen-1):,:]
+        else:
+            self.previous_feature = feature
+        return feature
+
+
     def forward(self, input, J_regressor=None):
         # input size NTF
         batch_size, seqlen, nc, h, w = input.shape
 
         feature = self.hmr.feature_extractor(input.reshape(-1, nc, h, w))
+
+
+
+        if(self.live_inference):
+            feature = self.generate_complete_feature_vector(feature)
+            seqlen = feature.shape[0]
+
 
         feature = feature.reshape(batch_size, seqlen, -1)
         feature = self.encoder(feature)

lib/utils/async_cam.py

@@ -0,0 +1,92 @@
+import cv2
+from threading import Thread
+import time
+import numpy as np
+
+# Install winCurses from pypi for only windows
+import curses
+
+
+class AsyncCamera(object):
+	def __init__(self,cam,display = False):
+
+		self.prev_frame_lis = []
+		self.frame = None
+		self.bbox = None
+		self.display_image = None
+		self.capture = cv2.VideoCapture(cam)
+		self.stop = False
+		self.display = display
+
+		self.thread = Thread(target=self.update, args=())
+		self.thread.daemon = True
+		self.thread.start()
+		if self.display:
+			self.start_displaying()
+
+
+	def check_key_press(self):
+		def key_press_wrapper(stdscr):
+			stdscr.nodelay(True)  
+			return stdscr.getch()==ord('q')
+		return curses.wrapper(key_press_wrapper)
+
+	def stop_cam(self):
+		self.capture.release()
+		self.stop = True
+
+
+	def show_frame(self):
+		image = np.copy(self.display_image)
+		if self.status and self.display_image is not None:
+			c_x,c_y,w,h = self.bbox
+			start_point = (int(c_x-w/2),int(c_y-h/2))
+			end_point = (int(c_x+w/2),int(c_y+h/2))
+			cv2.rectangle(image, start_point, end_point, (0, 0, 255) , 2) 
+			cv2.imshow('bbox', image)
+
+		key = cv2.waitKey(1)
+		if key == ord('q'):
+			self.stop_cam()
+			cv2.destroyAllWindows()
+
+
+
+	def update(self):
+		while True and not self.stop:
+			if self.capture.isOpened():
+				if(not self.display):
+					if(self.check_key_press()):
+						self.stop_cam()
+				(self.status, self.frame) = self.capture.read()
+				self.prev_frame_lis.append(self.frame)
+
+
+
+	def read(self):
+		if(self.frame is not None):
+			return True,self.prev_frame_lis
+		else:
+			return False,None
+
+	def del_frame_lis(self):
+		self.prev_frame_lis = []
+		self.status = False
+
+
+	def start_displaying(self):
+		def start_display_thread():
+			while True:
+				try:
+					self.show_frame()
+				except AttributeError:
+					pass
+		self.displayThread= Thread(target=start_display_thread, args=())
+		self.displayThread.daemon = True
+		self.displayThread.start()
+
+	def set_bounding_box(self,bbox):
+		self.bbox = bbox
+
+	def set_display_image(self,display_image):
+		self.display_image = display_image