add some code

tf-pose-estimation/inference_video.py

@@ -42,16 +42,19 @@ def addTopModelMobileNet(bottom_model,num_classes):
     parser.add_argument('--show-process', type=bool, default=False,
                         help='for debug purpose, if enabled, speed for inference is dropped.')
     parser.add_argument('--showBG', type=bool, default=True, help='False to show skeleton only.')
+    parser.add_argument('--num',type=int,default = 3)
+    parser.add_argument('--weight',type=str,default= 'weight.hdf5')
     args = parser.parse_args()
 
     logger.debug('initialization %s : %s' % (args.model, get_graph_path(args.model)))
     w, h = model_wh(args.resolution)
     e = TfPoseEstimator(get_graph_path(args.model), target_size=(w, h))
 
     mobilenet = MobileNet(include_top=False, input_shape=(224,224,3))
-    fc = addTopModelMobileNet(mobilenet,3)
+    output_num = args.num
+    fc = addTopModelMobileNet(mobilenet,output_num)
     mobile = Model(inputs=mobilenet.input,outputs=fc)
-    mobile.load_weights('weight.hdf5')
+    mobile.load_weights('./weight/'+args.weight)
 
     cap = cv2.VideoCapture(args.video)
 
@@ -70,7 +73,7 @@ def addTopModelMobileNet(bottom_model,num_classes):
 
         humans = e.inference(image,resize_to_default=(w>0 and h>0), upsample_size=4.0)
 
-        image,a,b = TfPoseEstimator.draw_humans(image,humans,mobile,imgcopy=False)
+        image,a,b = TfPoseEstimator.draw_humans(image,humans,mobile,output_num,imgcopy=False)
 
         cv2.putText(image,"FPS: %f" % (1.0 / (time.time() - fps_time)), (10,10), cv2.FONT_HERSHEY_SIMPLEX,0.5,(0,255,0),2)
 

tf-pose-estimation/run.py

@@ -61,11 +61,11 @@
         a.set_title('Result')
         image2 = cv2.cvtColor(image,cv2.COLOR_BGR2RGB)
         #plt.imshow(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
-        cv2.imwrite('../../data/'+name,image2)
-        for idx,img in enumerate(black_li):
-            cv2.imwrite('../../data/'+name.split('.')[0]+'_'+str(idx)+'.jpg',img)
-        for idx,img in enumerate(normal_li):
-            cv2.imwrite('../../data/'+name.split('.')[0]+'_normal'+str(idx)+'.jpg',cv2.cvtColor(img,cv2.COLOR_BGR2RGB))
+        cv2.imwrite('../../data/c_'+name,image2)
+        #for idx,img in enumerate(black_li):
+        #    cv2.imwrite('../../data/'+name.split('.')[0]+'_'+str(idx)+'.jpg',img)
+        #for idx,img in enumerate(normal_li):
+        #    cv2.imwrite('../../data/'+name.split('.')[0]+'_normal'+str(idx)+'.jpg',cv2.cvtColor(img,cv2.COLOR_BGR2RGB))
         print('../../data/'+name)
         bgimg = cv2.cvtColor(image.astype(np.uint8), cv2.COLOR_BGR2RGB)
         bgimg = cv2.resize(bgimg, (e.heatMat.shape[1], e.heatMat.shape[0]), interpolation=cv2.INTER_AREA)

tf-pose-estimation/test.py

@@ -0,0 +1,85 @@
+import argparse
+import logging
+import time
+
+import cv2
+import numpy as np
+
+from tf_pose.estimator import TfPoseEstimator
+from tf_pose.networks import get_graph_path, model_wh
+
+from tensorflow import keras
+from tensorflow.keras.models import Model
+from tensorflow.keras.applications import MobileNet
+from tensorflow.keras.layers import Dense,Activation,GlobalAveragePooling2D
+
+def addTopModelMobileNet(bottom_model, num_classes):
+    top_model = bottom_model.output
+    top_model = GlobalAveragePooling2D()(top_model)
+    top_model = Dense(1024,activation='relu')(top_model)
+    top_model = Dense(1024,activation='relu')(top_model)
+    top_model = Dense(512,activation='relu')(top_model)
+    top_model = Dense(num_classes,activation='softmax')(top_model)
+
+    return top_model
+
+
+logger = logging.getLogger('TfPoseEstimator-Video')
+logger.setLevel(logging.DEBUG)
+ch = logging.StreamHandler()
+ch.setLevel(logging.DEBUG)
+formatter = logging.Formatter('[%(asctime)s] [%(name)s] [%(levelname)s] %(message)s')
+ch.setFormatter(formatter)
+logger.addHandler(ch)
+
+fps_time = 0
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(description='tf-pose-estimation Video')
+    parser.add_argument('--video', type=str, default='../../data/video/crawl/IMG_4641.mp4')
+    parser.add_argument('--resolution', type=str, default='432x368', help='network input resolution. default=432x368')
+    parser.add_argument('--model', type=str, default='mobilenet_thin', help='cmu / mobilenet_thin / mobilenet_v2_large / mobilenet_v2_small')
+    parser.add_argument('--show-process', type=bool, default=False,
+                        help='for debug purpose, if enabled, speed for inference is dropped.')
+    parser.add_argument('--showBG', type=bool, default=True, help='False to show skeleton only.')
+    args = parser.parse_args()
+
+    mobilenet = MobileNet(include_top=False , input_shape=(224,224,3))
+    fc = addTopModelMobileNet(mobilenet,3)
+    mobile = Model(inputs=mobilenet.input,outputs=fc)
+    mobile.load_weights('weight.hdf5')
+
+    #logger.debug('initialization %s : %s' % (args.model, get_graph_path(args.model)))
+    w, h = model_wh(args.resolution)
+    print('hi')
+    e = TfPoseEstimator(get_graph_path(args.model), target_size=(w, h))
+    cap = cv2.VideoCapture(args.video)
+
+    if cap.isOpened() is False:
+        print("Error opening video stream or file")
+    num = 0
+    while cap.isOpened():
+        ret_val, image = cap.read()
+        #image = cv2.resize(image,(w,h))
+        print('hi')
+        humans = e.inference(image,resize_to_default=(w>0 and h>0 ), upsample_size=4.0)
+        print('hi')
+        if not args.showBG:
+            image = np.zeros(image.shape)
+        image,a,b = TfPoseEstimator.draw_humans(image, humans,mobile, imgcopy=False)
+
+        cv2.putText(image, "FPS: %f" % (1.0 / (time.time() - fps_time)), (10, 10),  cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
+
+        cv2.imwrite(f'../../data/inference/{num}.jpg',image)
+        print(num)
+        num+=1
+
+        #cv2.imshow('tf-pose-estimation result', image)
+        fps_time = time.time()
+        #if cv2.waitKey(1) == 27:
+        #    break
+
+    #cv2.destroyAllWindows()
+    print('finish!!')
+logger.debug('finished+')

tf-pose-estimation/tf_pose/estimator.py

@@ -412,7 +412,7 @@ def _quantize_img(npimg):
         return npimg_q
 
     @staticmethod
-    def draw_humans(npimg, humans,model, imgcopy=False):
+    def draw_humans(npimg, humans,model,output_num, imgcopy=False):
         if imgcopy:
             npimg = np.copy(npimg)
         image_h, image_w = npimg.shape[:2]
@@ -468,10 +468,15 @@ def draw_humans(npimg, humans,model, imgcopy=False):
 
            # black_li.append(crop)
             #normal_li.append(cv2.cvtColor(crop2,cv2.COLOR_BGR2RGB))
-            if output==2:
+            if output==0:
                 cv2.rectangle(npimg,(min_x-margin,max_y+margin),(max_x+margin,min_y-margin),(0,255,0),3)
-            else:
-                cv2.rectangle(npimg,(min_x-margin,max_y+margin),(max_x+margin,min_y-margin),(0,0,255),3)
+            elif output==1:
+                cv2.rectangle(npimg,(min_x-margin,max_y+margin),(max_x+margin,min_y-margin),(255,0,0),3)
+            elif output==2:
+                cv2.rectangle(npimg,(min_x-margin,max_y+margin),(max_x+margin,min_y-margin),(255,255,255),3)
+            if output_num==4:
+                if output==3:
+                    cv2.rectangle(npimg,(min_x-margin,max_y+margin),(max_x+margin,min_y-margin),(0,0,255),3)
             #print('humans body box',human.get_upper_body_box(image_w,image_h))
             #print('image',image_h,image_w)
         if no==len(humans):