Refactoring

app.py

@@ -2,7 +2,7 @@ import gradio as gr
 import os
 import wget
 import subprocess
-subprocess.call(['pip', 'install', 'git+https://github.com/facebookresearch/detectron2@main#subdirectory=projects/DensePose'])
+#subprocess.call(['pip', 'install', 'git+https://github.com/facebookresearch/detectron2@main#subdirectory=projects/DensePose'])
 from helpers.processor import TextureProcessor
 
 def image_processing(person_img, model_img):
@@ -12,12 +12,13 @@ def load_model(current_path):
     data_path = os.path.join(current_path, 'data')
     if not os.path.isdir(data_path):
         os.mkdir(data_path)
-        url = 'https://raw.githubusercontent.com/facebookresearch/detectron2/main/projects/DensePose/configs/densepose_rcnn_R_50_FPN_WC1M_s1x.yaml'
-        wget.download(url, os.path.join(data_path, 'config.yaml'))
-        url = 'https://dl.fbaipublicfiles.com/densepose/densepose_rcnn_R_50_FPN_WC1M_s1x/217144516/model_final_48a9d9.pkl'
-        wget.download(url, os.path.join(data_path, 'weights.pkl'))
-        url = 'https://raw.githubusercontent.com/facebookresearch/detectron2/main/projects/DensePose/configs/Base-DensePose-RCNN-FPN.yaml'
-        wget.download(url, os.path.join(data_path, 'Base-DensePose-RCNN-FPN.yaml'))
+        items_to_load = {
+            'config.yaml': 'https://raw.githubusercontent.com/facebookresearch/detectron2/main/projects/DensePose/configs/densepose_rcnn_R_50_FPN_WC1M_s1x.yaml',
+            'weights.pkl': 'https://dl.fbaipublicfiles.com/densepose/densepose_rcnn_R_50_FPN_WC1M_s1x/217144516/model_final_48a9d9.pkl',
+            'Base-DensePose-RCNN-FPN.yaml': 'https://raw.githubusercontent.com/facebookresearch/detectron2/main/projects/DensePose/configs/Base-DensePose-RCNN-FPN.yaml'
+        }
+        for filename, url in items_to_load.items():
+            wget.download(url, os.path.join(data_path, filename))
 
 current_path = os.getcwd()
 load_model(current_path)

helpers/processor.py

@@ -75,107 +75,107 @@ class TextureProcessor:
     return image_vis
 
   def parse_iuv(self, result):
-      i = result['pred_densepose'][0].labels.cpu().numpy().astype(float)
-      uv = (result['pred_densepose'][0].uv.cpu().numpy() * 255.0).astype(float)
-      iuv = np.stack((uv[1, :, :], uv[0, :, :], i))
-      iuv = np.transpose(iuv, (1, 2, 0))
-      return iuv
+    i = result['pred_densepose'][0].labels.cpu().numpy().astype(float)
+    uv = (result['pred_densepose'][0].uv.cpu().numpy() * 255.0).astype(float)
+    iuv = np.stack((uv[1, :, :], uv[0, :, :], i))
+    iuv = np.transpose(iuv, (1, 2, 0))
+    return iuv
 
   def parse_bbox(self, result):
-      return result['pred_boxes_XYXY'][0].cpu().numpy()
+    return result['pred_boxes_XYXY'][0].cpu().numpy()
 
   def interpolate_tex(self, tex):
-      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)
-      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 = actual_part * dilated_mask
-
-      return actual_part
+    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)
+    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 = 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)
-      return texture
+    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)
+    return texture
 
   def get_texture(self, im, iuv, bbox, tex_part_size=200):
-      im = im.transpose(2, 1, 0) / 255
-      image_w, image_h = im.shape[1], im.shape[2]
-      bbox[2] = bbox[2] - bbox[0]
-      bbox[3] = bbox[3] - bbox[1]
-      x, y, w, h = [int(v) for v in bbox]
-      bg = np.zeros((image_h, image_w, 3))
-      bg[y:y + h, x:x + w, :] = iuv
-      iuv = bg
-      iuv = iuv.transpose((2, 1, 0))
-      i, u, v = iuv[2], iuv[1], iuv[0]
+    im = im.transpose(2, 1, 0) / 255
+    image_w, image_h = im.shape[1], im.shape[2]
+    bbox[2] = bbox[2] - bbox[0]
+    bbox[3] = bbox[3] - bbox[1]
+    x, y, w, h = [int(v) for v in bbox]
+    bg = np.zeros((image_h, image_w, 3))
+    bg[y:y + h, x:x + w, :] = iuv
+    iuv = bg
+    iuv = iuv.transpose((2, 1, 0))
+    i, u, v = iuv[2], iuv[1], iuv[0]
 
-      n_parts = 22
-      texture = np.zeros((n_parts, 3, tex_part_size, tex_part_size))
+    n_parts = 22
+    texture = np.zeros((n_parts, 3, tex_part_size, tex_part_size))
 
-      for part_id in range(1, n_parts + 1):
-          generated = np.zeros((3, tex_part_size, tex_part_size))
+    for part_id in range(1, n_parts + 1):
+        generated = np.zeros((3, tex_part_size, tex_part_size))
 
-          x, y = u[i == part_id], v[i == part_id]
+        x, y = u[i == part_id], v[i == part_id]
 
-          tex_u_coo = (x * (tex_part_size - 1) / 255).astype(int)
-          tex_v_coo = (y * (tex_part_size - 1) / 255).astype(int)
+        tex_u_coo = (x * (tex_part_size - 1) / 255).astype(int)
+        tex_v_coo = (y * (tex_part_size - 1) / 255).astype(int)
 
-          tex_u_coo = np.clip(tex_u_coo, 0, tex_part_size - 1)
-          tex_v_coo = np.clip(tex_v_coo, 0, tex_part_size - 1)
+        tex_u_coo = np.clip(tex_u_coo, 0, tex_part_size - 1)
+        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]
+        for channel in range(3):
+            generated[channel][tex_v_coo,
+                tex_u_coo] = im[channel][i == part_id]
 
-          if np.sum(generated) > 0:
-              generated = self.interpolate_tex(generated)
+        if np.sum(generated) > 0:
+            generated = self.interpolate_tex(generated)
 
-          texture[part_id - 1] = generated[:, ::-1, :]
+        texture[part_id - 1] = generated[:, ::-1, :]
 
-      tex_concat = np.zeros((24, tex_part_size, tex_part_size, 3))
-      for i in range(texture.shape[0]):
-          tex_concat[i] = texture[i].transpose(2, 1, 0)
-      tex = self.concat_textures(tex_concat)
+    tex_concat = np.zeros((24, tex_part_size, tex_part_size, 3))
+    for i in range(texture.shape[0]):
+        tex_concat[i] = texture[i].transpose(2, 1, 0)
+    tex = self.concat_textures(tex_concat)
 
-      return tex
+    return tex
 
   def create_iuv(self, results, image):
-      iuv = self.parse_iuv(results)
-      bbox = self.parse_bbox(results)
-      uv_texture = self.get_texture(image, iuv, bbox)
-      uv_texture = uv_texture.transpose([1, 0, 2])
-      return uv_texture
+    iuv = self.parse_iuv(results)
+    bbox = self.parse_bbox(results)
+    uv_texture = self.get_texture(image, iuv, bbox)
+    uv_texture = uv_texture.transpose([1, 0, 2])
+    return uv_texture
 
   def get_config(self, config_fpath, model_fpath):
-      cfg = get_cfg()
-      add_densepose_config(cfg)
-      cfg.merge_from_file(config_fpath)
-      cfg.MODEL.WEIGHTS = model_fpath
-      cfg.MODEL.DEVICE = 'cpu'
-      cfg.freeze()
-      return cfg
+    cfg = get_cfg()
+    add_densepose_config(cfg)
+    cfg.merge_from_file(config_fpath)
+    cfg.MODEL.WEIGHTS = model_fpath
+    cfg.MODEL.DEVICE = 'cpu'
+    cfg.freeze()
+    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']
+    context = {'results': []}
+    with torch.no_grad():
+        outputs = self.predictor(image)['instances']
+        self.execute_on_outputs(context, outputs)
+    return context['results']
 
   def execute_on_outputs(self, context: Dict[str, Any], outputs: Instances):
     result = {}