Merge pull request #1118 from Laicheng0830/add_pose

Add 3D human poses

tensorlayer/app/__init__.py

@@ -2,4 +2,5 @@
 # -*- coding: utf-8 -*-
 
 from .computer_vision_object_detection import *
+from .human_pose_estimation import *
 from .computer_vision import *

tensorlayer/app/computer_vision.py

@@ -2,18 +2,48 @@
 # -*- coding: utf-8 -*-
 
 from tensorlayer.app import YOLOv4, get_anchors, decode, filter_boxes
+from tensorlayer.app import CGCNN
 import numpy as np
 import tensorflow as tf
 from tensorlayer import logging
 import cv2
 
 
 class object_detection(object):
+    """Model encapsulation.
+
+    Parameters
+    ----------
+    model_name : str
+        Choose the model to inference.
+
+    Methods
+    ---------
+    __init__()
+        Initializing the model.
+    __call__()
+        (1)Formatted input and output. (2)Inference model.
+    list()
+        Abstract method. Return available a list of model_name.
+
+    Examples
+    ---------
+    Object Detection detection MSCOCO with YOLOv4, see `tutorial_object_detection_yolov4.py
+    <https://github.com/tensorlayer/tensorlayer/blob/master/example/app_tutorials/tutorial_object_detection_yolov4.py>`__
+    With TensorLayer
+
+    >>> # get the whole model
+    >>> net = tl.app.computer_vision.object_detection('yolo4-mscoco')
+    >>> # use for inferencing
+    >>> output = net(img)
+    """
 
     def __init__(self, model_name='yolo4-mscoco'):
         self.model_name = model_name
         if self.model_name == 'yolo4-mscoco':
             self.model = YOLOv4(NUM_CLASS=80, pretrained=True)
+        elif self.model_name == 'lcn':
+            self.model = CGCNN(pretrained=True)
         else:
             raise ("The model does not support.")
 
@@ -23,6 +53,8 @@ def __call__(self, input_data):
             feature_maps = self.model(batch_data, is_train=False)
             pred_bbox = yolo4_output_processing(feature_maps)
             output = result_to_json(input_data, pred_bbox)
+        elif self.model_name == 'lcn':
+            output = self.model(input_data)
         else:
             raise NotImplementedError
 
@@ -35,7 +67,7 @@ def __repr__(self):
 
     @property
     def list(self):
-        logging.info("The model name list: yolov4-mscoco")
+        logging.info("The model name list: 'yolov4-mscoco', 'lcn'")
 
 
 def yolo4_input_processing(original_image):

tensorlayer/app/computer_vision_object_detection/yolov4.py

@@ -1,6 +1,6 @@
 #! /usr/bin/python
 # -*- coding: utf-8 -*-
-"""YOLOv4 for MS COCO.
+"""YOLOv4 for MS-COCO.
 
 # Reference:
 - [tensorflow-yolov4-tflite](
@@ -139,6 +139,29 @@ def cspdarknet53(input_data=None):
 
 
 def YOLOv4(NUM_CLASS, pretrained=False):
+    """Pre-trained YOLOv4 model.
+
+    Parameters
+    ------------
+    NUM_CLASS : int
+        Number of classes in final prediction.
+    pretrained : boolean
+        Whether to load pretrained weights. Default False.
+
+    Examples
+    ---------
+    Object Detection with YOLOv4, see `computer_vision.py
+    <https://github.com/tensorlayer/tensorlayer/blob/master/tensorlayer/app/computer_vision.py>`__
+    With TensorLayer
+
+    >>> # get the whole model, without pre-trained YOLOv4 parameters
+    >>> yolov4 = tl.app.YOLOv4(NUM_CLASS=80, pretrained=False)
+    >>> # get the whole model, restore pre-trained YOLOv4 parameters
+    >>> yolov4 = tl.app.YOLOv4(NUM_CLASS=80, pretrained=True)
+    >>> # use for inferencing
+    >>> output = yolov4(img, is_train=False)
+
+    """
 
     input_layer = Input([None, INPUT_SIZE, INPUT_SIZE, 3])
     route_1, route_2, conv = cspdarknet53(input_layer)