Spaces:
Runtime error
Runtime error
Processor refactoring
Browse files- helpers/processor.py +13 -26
helpers/processor.py
CHANGED
@@ -2,7 +2,6 @@ import cv2
|
|
2 |
import imageio
|
3 |
import numpy as np
|
4 |
import torch
|
5 |
-
from typing import Any, Dict
|
6 |
import io
|
7 |
|
8 |
from detectron2.config import get_cfg
|
@@ -10,10 +9,6 @@ from detectron2.engine.defaults import DefaultPredictor
|
|
10 |
from detectron2.structures.instances import Instances
|
11 |
|
12 |
from densepose import add_densepose_config
|
13 |
-
from densepose.structures import (
|
14 |
-
DensePoseChartPredictorOutput,
|
15 |
-
DensePoseEmbeddingPredictorOutput
|
16 |
-
)
|
17 |
from densepose.vis.base import CompoundVisualizer
|
18 |
from densepose.vis.densepose_outputs_vertex import get_texture_atlases
|
19 |
from densepose.vis.densepose_results_textures import (
|
@@ -22,7 +17,6 @@ from densepose.vis.densepose_results_textures import (
|
|
22 |
from densepose.vis.extractor import (
|
23 |
CompoundExtractor,
|
24 |
create_extractor,
|
25 |
-
DensePoseOutputsExtractor,
|
26 |
DensePoseResultExtractor
|
27 |
)
|
28 |
|
@@ -34,8 +28,7 @@ class TextureProcessor:
|
|
34 |
|
35 |
def process_texture(self, image):
|
36 |
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
|
37 |
-
output = self.execute(image)
|
38 |
-
|
39 |
if 'pred_densepose' in output:
|
40 |
texture = self.create_iuv(output, image)
|
41 |
atlas_texture_bytes = io.BytesIO()
|
@@ -90,26 +83,22 @@ class TextureProcessor:
|
|
90 |
radius_increase = 10
|
91 |
kernel = np.ones((radius_increase, radius_increase), np.uint8)
|
92 |
dilated_mask = cv2.dilate(valid_mask, kernel, iterations=1)
|
|
|
93 |
invalid_region = 1 - valid_mask
|
94 |
actual_part_max = tex.max()
|
95 |
actual_part_min = tex.min()
|
96 |
-
actual_part_uint = np.array(
|
97 |
-
|
98 |
-
actual_part_uint = cv2.inpaint(actual_part_uint.transpose((1, 2, 0)), invalid_region, 1,
|
99 |
-
|
100 |
-
actual_part = (actual_part_uint / 255.0) *
|
101 |
-
(actual_part_max - actual_part_min) + actual_part_min
|
102 |
actual_part = actual_part * dilated_mask
|
103 |
-
|
104 |
return actual_part
|
105 |
|
106 |
def concat_textures(self, array):
|
107 |
-
|
108 |
-
|
109 |
-
tmp = array[6 * i]
|
110 |
-
for j in range(6 * i + 1, 6 * i + 6):
|
111 |
-
tmp = np.concatenate((tmp, array[j]), axis=1)
|
112 |
-
texture = tmp if len(texture) == 0 else np.concatenate((texture, tmp), axis=0)
|
113 |
return texture
|
114 |
|
115 |
def get_texture(self, im, iuv, bbox, tex_part_size=200):
|
@@ -170,13 +159,11 @@ class TextureProcessor:
|
|
170 |
return cfg
|
171 |
|
172 |
def execute(self, image):
|
173 |
-
context = {'results': []}
|
174 |
with torch.no_grad():
|
175 |
outputs = self.predictor(image)['instances']
|
176 |
-
self.execute_on_outputs(
|
177 |
-
return context['results']
|
178 |
|
179 |
-
def execute_on_outputs(self,
|
180 |
result = {}
|
181 |
if outputs.has('scores'):
|
182 |
result['scores'] = outputs.get('scores').cpu()
|
@@ -184,4 +171,4 @@ class TextureProcessor:
|
|
184 |
result['pred_boxes_XYXY'] = outputs.get('pred_boxes').tensor.cpu()
|
185 |
if outputs.has('pred_densepose'):
|
186 |
result['pred_densepose'] = self.extractor(outputs)[0]
|
187 |
-
|
|
|
2 |
import imageio
|
3 |
import numpy as np
|
4 |
import torch
|
|
|
5 |
import io
|
6 |
|
7 |
from detectron2.config import get_cfg
|
|
|
9 |
from detectron2.structures.instances import Instances
|
10 |
|
11 |
from densepose import add_densepose_config
|
|
|
|
|
|
|
|
|
12 |
from densepose.vis.base import CompoundVisualizer
|
13 |
from densepose.vis.densepose_outputs_vertex import get_texture_atlases
|
14 |
from densepose.vis.densepose_results_textures import (
|
|
|
17 |
from densepose.vis.extractor import (
|
18 |
CompoundExtractor,
|
19 |
create_extractor,
|
|
|
20 |
DensePoseResultExtractor
|
21 |
)
|
22 |
|
|
|
28 |
|
29 |
def process_texture(self, image):
|
30 |
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
|
31 |
+
output = self.execute(image)
|
|
|
32 |
if 'pred_densepose' in output:
|
33 |
texture = self.create_iuv(output, image)
|
34 |
atlas_texture_bytes = io.BytesIO()
|
|
|
83 |
radius_increase = 10
|
84 |
kernel = np.ones((radius_increase, radius_increase), np.uint8)
|
85 |
dilated_mask = cv2.dilate(valid_mask, kernel, iterations=1)
|
86 |
+
|
87 |
invalid_region = 1 - valid_mask
|
88 |
actual_part_max = tex.max()
|
89 |
actual_part_min = tex.min()
|
90 |
+
actual_part_uint = np.array((tex - actual_part_min) / (actual_part_max - actual_part_min) * 255, dtype='uint8')
|
91 |
+
|
92 |
+
actual_part_uint = cv2.inpaint(actual_part_uint.transpose((1, 2, 0)), invalid_region, 1, cv2.INPAINT_TELEA).transpose((2, 0, 1))
|
93 |
+
|
94 |
+
actual_part = (actual_part_uint / 255.0) * (actual_part_max - actual_part_min) + actual_part_min
|
|
|
95 |
actual_part = actual_part * dilated_mask
|
96 |
+
|
97 |
return actual_part
|
98 |
|
99 |
def concat_textures(self, array):
|
100 |
+
texture_rows = [np.concatenate(array[i:i+6], axis=1) for i in range(0, 24, 6)]
|
101 |
+
texture = np.concatenate(texture_rows, axis=0)
|
|
|
|
|
|
|
|
|
102 |
return texture
|
103 |
|
104 |
def get_texture(self, im, iuv, bbox, tex_part_size=200):
|
|
|
159 |
return cfg
|
160 |
|
161 |
def execute(self, image):
|
|
|
162 |
with torch.no_grad():
|
163 |
outputs = self.predictor(image)['instances']
|
164 |
+
return self.execute_on_outputs(outputs)
|
|
|
165 |
|
166 |
+
def execute_on_outputs(self, outputs: Instances):
|
167 |
result = {}
|
168 |
if outputs.has('scores'):
|
169 |
result['scores'] = outputs.get('scores').cpu()
|
|
|
171 |
result['pred_boxes_XYXY'] = outputs.get('pred_boxes').tensor.cpu()
|
172 |
if outputs.has('pred_densepose'):
|
173 |
result['pred_densepose'] = self.extractor(outputs)[0]
|
174 |
+
return result
|