Skip to content

Commit 3965000

Browse files
authored
Merge pull request #1118 from Laicheng0830/add_pose
Add 3D human poses
2 parents 39c3159 + fa8c66b commit 3965000

6 files changed

Lines changed: 445 additions & 2 deletions

File tree

tensorlayer/app/__init__.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,4 +2,5 @@
22
# -*- coding: utf-8 -*-
33

44
from .computer_vision_object_detection import *
5+
from .human_pose_estimation import *
56
from .computer_vision import *

tensorlayer/app/computer_vision.py

Lines changed: 33 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,18 +2,48 @@
22
# -*- coding: utf-8 -*-
33

44
from tensorlayer.app import YOLOv4, get_anchors, decode, filter_boxes
5+
from tensorlayer.app import CGCNN
56
import numpy as np
67
import tensorflow as tf
78
from tensorlayer import logging
89
import cv2
910

1011

1112
class object_detection(object):
13+
"""Model encapsulation.
14+
15+
Parameters
16+
----------
17+
model_name : str
18+
Choose the model to inference.
19+
20+
Methods
21+
---------
22+
__init__()
23+
Initializing the model.
24+
__call__()
25+
(1)Formatted input and output. (2)Inference model.
26+
list()
27+
Abstract method. Return available a list of model_name.
28+
29+
Examples
30+
---------
31+
Object Detection detection MSCOCO with YOLOv4, see `tutorial_object_detection_yolov4.py
32+
<https://github.com/tensorlayer/tensorlayer/blob/master/example/app_tutorials/tutorial_object_detection_yolov4.py>`__
33+
With TensorLayer
34+
35+
>>> # get the whole model
36+
>>> net = tl.app.computer_vision.object_detection('yolo4-mscoco')
37+
>>> # use for inferencing
38+
>>> output = net(img)
39+
"""
1240

1341
def __init__(self, model_name='yolo4-mscoco'):
1442
self.model_name = model_name
1543
if self.model_name == 'yolo4-mscoco':
1644
self.model = YOLOv4(NUM_CLASS=80, pretrained=True)
45+
elif self.model_name == 'lcn':
46+
self.model = CGCNN(pretrained=True)
1747
else:
1848
raise ("The model does not support.")
1949

@@ -23,6 +53,8 @@ def __call__(self, input_data):
2353
feature_maps = self.model(batch_data, is_train=False)
2454
pred_bbox = yolo4_output_processing(feature_maps)
2555
output = result_to_json(input_data, pred_bbox)
56+
elif self.model_name == 'lcn':
57+
output = self.model(input_data)
2658
else:
2759
raise NotImplementedError
2860

@@ -35,7 +67,7 @@ def __repr__(self):
3567

3668
@property
3769
def list(self):
38-
logging.info("The model name list: yolov4-mscoco")
70+
logging.info("The model name list: 'yolov4-mscoco', 'lcn'")
3971

4072

4173
def yolo4_input_processing(original_image):

tensorlayer/app/computer_vision_object_detection/yolov4.py

Lines changed: 24 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
#! /usr/bin/python
22
# -*- coding: utf-8 -*-
3-
"""YOLOv4 for MS COCO.
3+
"""YOLOv4 for MS-COCO.
44
55
# Reference:
66
- [tensorflow-yolov4-tflite](
@@ -139,6 +139,29 @@ def cspdarknet53(input_data=None):
139139

140140

141141
def YOLOv4(NUM_CLASS, pretrained=False):
142+
"""Pre-trained YOLOv4 model.
143+
144+
Parameters
145+
------------
146+
NUM_CLASS : int
147+
Number of classes in final prediction.
148+
pretrained : boolean
149+
Whether to load pretrained weights. Default False.
150+
151+
Examples
152+
---------
153+
Object Detection with YOLOv4, see `computer_vision.py
154+
<https://github.com/tensorlayer/tensorlayer/blob/master/tensorlayer/app/computer_vision.py>`__
155+
With TensorLayer
156+
157+
>>> # get the whole model, without pre-trained YOLOv4 parameters
158+
>>> yolov4 = tl.app.YOLOv4(NUM_CLASS=80, pretrained=False)
159+
>>> # get the whole model, restore pre-trained YOLOv4 parameters
160+
>>> yolov4 = tl.app.YOLOv4(NUM_CLASS=80, pretrained=True)
161+
>>> # use for inferencing
162+
>>> output = yolov4(img, is_train=False)
163+
164+
"""
142165

143166
input_layer = Input([None, INPUT_SIZE, INPUT_SIZE, 3])
144167
route_1, route_2, conv = cspdarknet53(input_layer)

0 commit comments

Comments
 (0)