Initial commit

.gitignore

@@ -0,0 +1,3 @@
+data/
+flagged/
+texture.png

app.py

@@ -0,0 +1,42 @@
+import gradio as gr
+import os
+import wget
+from helpers.processor import TextureProcessor
+
+def image_processing(person_img, model_img):
+    return texture_processor.extract(person_img, model_img)
+
+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'))
+
+current_path = os.getcwd()
+load_model(current_path)
+densepose_config = os.path.join(current_path, 'data', 'config.yaml')
+densepose_weights = os.path.join(current_path, 'data', 'weights.pkl')
+
+texture_processor = TextureProcessor(densepose_config, densepose_weights)
+
+inputs = [
+    gr.inputs.Image(label="Person Image", type='numpy'),
+    gr.inputs.Image(label="Model Image (with clothes)", type='numpy')
+]
+
+outputs = gr.outputs.Image(label="Result Image", type='numpy')
+
+title = "JustClothify"
+description = "Upload an image of a person and image of model with clothes, and the model will generate image of the person, wearing this clothing."
+
+gr.Interface(
+    fn=image_processing, 
+    inputs=inputs, 
+    outputs=outputs, 
+    title=title, 
+    description=description).launch()

helpers/processor.py

@@ -0,0 +1,196 @@
+import cv2
+import imageio
+import numpy as np
+import torch
+from typing import Any, Dict
+
+from detectron2.config import get_cfg
+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 (
+    DensePoseResultsVisualizerWithTexture as dp_iuv_texture,
+    get_texture_atlas
+)
+from densepose.vis.extractor import (
+    CompoundExtractor,
+    create_extractor,
+    DensePoseOutputsExtractor,
+    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):
+    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)
+    else:
+      print(f'Одежда не распознана')
+
+  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)
+
+  def overlay_texture(self, texture_name, original_image):
+    texture_atlas = get_texture_atlas(texture_name)
+    texture_atlases_dict = get_texture_atlases(None)
+    vis = dp_iuv_texture(
+        cfg=self.config,
+        texture_atlas=texture_atlas,
+        texture_atlases_dict=texture_atlases_dict
+    )
+
+    visualizers = [vis]
+    extractor = create_extractor(vis)
+    extractors = [extractor]
+
+    visualizer = CompoundVisualizer(visualizers)
+    extractor = CompoundExtractor(extractors)
+
+    #img = read_image(original_image, format="BGR")
+    with torch.no_grad():
+      outputs = self.predictor(original_image)["instances"]
+    
+    image = cv2.cvtColor(original_image, cv2.COLOR_BGR2GRAY)
+    image = np.tile(image[:, :, np.newaxis], [1, 1, 3])
+    data = extractor(outputs)
+    image_vis = visualizer.visualize(image, data)
+
+    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
+
+  def parse_bbox(self, result):
+      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()
+      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
+      # only use dilated part
+      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
+
+  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]
+
+      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))
+
+          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 = 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]
+
+          if np.sum(generated) > 0:
+              generated = self.interpolate_tex(generated)
+
+          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)
+
+      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
+
+  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
+
+  def execute(self, image):
+      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 = {}
+    if outputs.has("scores"):
+        result["scores"] = outputs.get("scores").cpu()
+    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]
+    context["results"].append(result)

requirements.txt

@@ -0,0 +1,2 @@
+git+https://github.com/facebookresearch/detectron2@main#subdirectory=projects/DensePose
+wget