Оптимизация

helpers/processor.py

@@ -3,6 +3,7 @@ import imageio
 import numpy as np
 import torch
 from typing import Any, Dict
+import io
 
 from detectron2.config import get_cfg
 from detectron2.engine.defaults import DefaultPredictor
@@ -16,8 +17,7 @@ from densepose.structures import (
 from densepose.vis.base import CompoundVisualizer
 from densepose.vis.densepose_outputs_vertex import get_texture_atlases
 from densepose.vis.densepose_results_textures import (
-    DensePoseResultsVisualizerWithTexture as dp_iuv_texture,
-    get_texture_atlas
+    DensePoseResultsVisualizerWithTexture as dp_iuv_texture
 )
 from densepose.vis.extractor import (
     CompoundExtractor,
@@ -26,30 +26,32 @@ from densepose.vis.extractor import (
     DensePoseResultExtractor
 )
 
-from detectron2.data.detection_utils import read_image
-
 class TextureProcessor:
   def __init__(self, config, weights):
     self.config = self.get_config(config, weights)
     self.predictor = DefaultPredictor(self.config)
   
-  def process_texture(self, image, output_filename):
+  def process_texture(self, image):
     image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
     output = self.execute(image)[0]
 
     if 'pred_densepose' in output:
         texture = self.create_iuv(output, image)
-        imageio.imwrite(output_filename, texture)
+        atlas_texture_bytes = io.BytesIO()
+        imageio.imwrite(atlas_texture_bytes, texture, format='PNG')
+        texture_atlas_array = np.frombuffer(atlas_texture_bytes.getvalue(), dtype=np.uint8)
+        texture_atlas = cv2.imdecode(texture_atlas_array, cv2.IMREAD_COLOR)
+        texture_atlas = cv2.cvtColor(texture_atlas, cv2.COLOR_BGR2RGB)
+        return texture_atlas
     else:
-      print(f'Одежда не распознана')
+      raise Exception('Clothes not found')
 
   def extract(self, person_img, model_img):
-    texture_filename = 'texture.png'
-    self.process_texture(model_img, texture_filename)
-    return self.overlay_texture(texture_filename, person_img)
+    texture_atlas = self.process_texture(model_img)
+    return self.overlay_texture(texture_atlas, person_img)
 
-  def overlay_texture(self, texture_name, original_image):
-    texture_atlas = get_texture_atlas(texture_name)
+  def overlay_texture(self, texture_atlas, original_image):
+    texture_atlas[:, :, :3] = texture_atlas[:, :, 2::-1]
     texture_atlases_dict = get_texture_atlases(None)
     vis = dp_iuv_texture(
         cfg=self.config,
@@ -57,14 +59,11 @@ class TextureProcessor:
         texture_atlases_dict=texture_atlases_dict
     )
 
-    visualizers = [vis]
     extractor = create_extractor(vis)
-    extractors = [extractor]
 
-    visualizer = CompoundVisualizer(visualizers)
-    extractor = CompoundExtractor(extractors)
+    visualizer = CompoundVisualizer([vis])
+    extractor = CompoundExtractor([extractor])
 
-    #img = read_image(original_image, format="BGR")
     with torch.no_grad():
       outputs = self.predictor(original_image)["instances"]
     
@@ -86,12 +85,10 @@ class TextureProcessor:
       return result["pred_boxes_XYXY"][0].cpu().numpy()
 
   def interpolate_tex(self, tex):
-      # code is adopted from https://github.com/facebookresearch/DensePose/issues/68
       valid_mask = np.array((tex.sum(0) != 0) * 1, dtype='uint8')
       radius_increase = 10
       kernel = np.ones((radius_increase, radius_increase), np.uint8)
       dilated_mask = cv2.dilate(valid_mask, kernel, iterations=1)
-      region_to_fill = dilated_mask - valid_mask
       invalid_region = 1 - valid_mask
       actual_part_max = tex.max()
       actual_part_min = tex.min()
@@ -101,7 +98,6 @@ class TextureProcessor:
                                     cv2.INPAINT_TELEA).transpose((2, 0, 1))
       actual_part = (actual_part_uint / 255.0) * \
           (actual_part_max - actual_part_min) + actual_part_min
-      # only use dilated part
       actual_part = actual_part * dilated_mask
 
       return actual_part