Optimizations

helpers/processor.py

@@ -30,6 +30,7 @@ class TextureProcessor:
   def __init__(self, config, weights):
     self.config = self.get_config(config, weights)
     self.predictor = DefaultPredictor(self.config)
+    self.extractor = DensePoseResultExtractor()
   
   def process_texture(self, image):
     image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
@@ -108,8 +109,7 @@ class TextureProcessor:
         tmp = array[6 * i]
         for j in range(6 * i + 1, 6 * i + 6):
             tmp = np.concatenate((tmp, array[j]), axis=1)
-        texture = tmp if len(texture) == 0 else np.concatenate(
-            (texture, tmp), axis=0)
+        texture = tmp if len(texture) == 0 else np.concatenate((texture, tmp), axis=0)
     return texture
 
   def get_texture(self, im, iuv, bbox, tex_part_size=200):
@@ -139,8 +139,7 @@ class TextureProcessor:
         tex_v_coo = np.clip(tex_v_coo, 0, tex_part_size - 1)
 
         for channel in range(3):
-            generated[channel][tex_v_coo,
-                tex_u_coo] = im[channel][i == part_id]
+            generated[channel][tex_v_coo, tex_u_coo] = im[channel][i == part_id]
 
         if np.sum(generated) > 0:
             generated = self.interpolate_tex(generated)
@@ -184,9 +183,5 @@ class TextureProcessor:
     if outputs.has('pred_boxes'):
         result['pred_boxes_XYXY'] = outputs.get('pred_boxes').tensor.cpu()
         if outputs.has('pred_densepose'):
-            if isinstance(outputs.pred_densepose, DensePoseChartPredictorOutput):
-                extractor = DensePoseResultExtractor()
-            elif isinstance(outputs.pred_densepose, DensePoseEmbeddingPredictorOutput):
-                extractor = DensePoseOutputsExtractor()
-            result['pred_densepose'] = extractor(outputs)[0]
+            result['pred_densepose'] = self.extractor(outputs)[0]
     context['results'].append(result)