Processor refactoring

helpers/processor.py

@@ -2,7 +2,6 @@ import cv2
 import imageio
 import numpy as np
 import torch
-from typing import Any, Dict
 import io
 
 from detectron2.config import get_cfg
@@ -10,10 +9,6 @@ from detectron2.engine.defaults import DefaultPredictor
 from detectron2.structures.instances import Instances
 
 from densepose import add_densepose_config
-from densepose.structures import (
-    DensePoseChartPredictorOutput,
-    DensePoseEmbeddingPredictorOutput
-)
 from densepose.vis.base import CompoundVisualizer
 from densepose.vis.densepose_outputs_vertex import get_texture_atlases
 from densepose.vis.densepose_results_textures import (
@@ -22,7 +17,6 @@ from densepose.vis.densepose_results_textures import (
 from densepose.vis.extractor import (
     CompoundExtractor,
     create_extractor,
-    DensePoseOutputsExtractor,
     DensePoseResultExtractor
 )
 
@@ -34,8 +28,7 @@ class TextureProcessor:
   
   def process_texture(self, image):
     image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
-    output = self.execute(image)[0]
-
+    output = self.execute(image)
     if 'pred_densepose' in output:
         texture = self.create_iuv(output, image)
         atlas_texture_bytes = io.BytesIO()
@@ -90,26 +83,22 @@ class TextureProcessor:
     radius_increase = 10
     kernel = np.ones((radius_increase, radius_increase), np.uint8)
     dilated_mask = cv2.dilate(valid_mask, kernel, iterations=1)
+    
     invalid_region = 1 - valid_mask
     actual_part_max = tex.max()
     actual_part_min = tex.min()
-    actual_part_uint = np.array(
-        (tex - actual_part_min) / (actual_part_max - actual_part_min) * 255, dtype='uint8')
-    actual_part_uint = cv2.inpaint(actual_part_uint.transpose((1, 2, 0)), invalid_region, 1,
-        cv2.INPAINT_TELEA).transpose((2, 0, 1))
-    actual_part = (actual_part_uint / 255.0) * \
-        (actual_part_max - actual_part_min) + actual_part_min
+    actual_part_uint = np.array((tex - actual_part_min) / (actual_part_max - actual_part_min) * 255, dtype='uint8')
+    
+    actual_part_uint = cv2.inpaint(actual_part_uint.transpose((1, 2, 0)), invalid_region, 1, cv2.INPAINT_TELEA).transpose((2, 0, 1))
+    
+    actual_part = (actual_part_uint / 255.0) * (actual_part_max - actual_part_min) + actual_part_min
     actual_part = actual_part * dilated_mask
-
+    
     return actual_part
 
   def concat_textures(self, array):
-    texture = []
-    for i in range(4):
-        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_rows = [np.concatenate(array[i:i+6], axis=1) for i in range(0, 24, 6)]
+    texture = np.concatenate(texture_rows, axis=0)
     return texture
 
   def get_texture(self, im, iuv, bbox, tex_part_size=200):
@@ -170,13 +159,11 @@ class TextureProcessor:
     return cfg
 
   def execute(self, image):
-    context = {'results': []}
     with torch.no_grad():
         outputs = self.predictor(image)['instances']
-        self.execute_on_outputs(context, outputs)
-    return context['results']
+        return self.execute_on_outputs(outputs)
 
-  def execute_on_outputs(self, context: Dict[str, Any], outputs: Instances):
+  def execute_on_outputs(self, outputs: Instances):
     result = {}
     if outputs.has('scores'):
         result['scores'] = outputs.get('scores').cpu()
@@ -184,4 +171,4 @@ class TextureProcessor:
         result['pred_boxes_XYXY'] = outputs.get('pred_boxes').tensor.cpu()
         if outputs.has('pred_densepose'):
             result['pred_densepose'] = self.extractor(outputs)[0]
-    context['results'].append(result)
+    return result