Refactoring

app.py

@@ -1,27 +1,29 @@
-import gradio as gr
 import os
-import wget
 import subprocess
-subprocess.call(['pip', 'install', 'git+https://github.com/facebookresearch/detectron2@main#subdirectory=projects/DensePose'])
+from typing import Dict
+import json
+import numpy as np
+import wget
+import gradio as gr
 from helpers.processor import TextureProcessor
 
-def image_processing(person_img, model_img):
+#subprocess.call(['pip', 'install', 'git+https://github.com/facebookresearch/detectron2@main#subdirectory=projects/DensePose'])
+
+def image_processing(person_img: np.ndarray, model_img: np.ndarray) -> np.ndarray:
     return texture_processor.extract(person_img, model_img)
 
-def load_model(current_path):
+def load_model(current_path: str, config: Dict) -> None:
     data_path = os.path.join(current_path, 'data')
     if not os.path.isdir(data_path):
         os.mkdir(data_path)
-        loader_dict = {
-            '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 loader_dict.items():
+        for filename, url in config.items():
             wget.download(url, os.path.join(data_path, filename))
 
+with open("config.json", "r") as f:
+    config = json.load(f)
+
 current_path = os.getcwd()
-load_model(current_path)
+load_model(current_path, config)
 densepose_config = os.path.join(current_path, 'data', 'config.yaml')
 densepose_weights = os.path.join(current_path, 'data', 'weights.pkl')
 
@@ -29,7 +31,7 @@ texture_processor = TextureProcessor(densepose_config, densepose_weights)
 
 title = '<h1 style="text-align:center">JustClothify</h1>'
 
-with gr.Blocks(theme='soft', title='AnswerMate') as blocks:
+with gr.Blocks(theme='soft', title='JustClothify') as blocks:
     gr.HTML(title)
     gr.Markdown('Upload an image of a person and an image of a model with clothes, the system will generate an image of a person wearing these clothes.')
     with gr.Row():

config.json

@@ -0,0 +1,5 @@
+{
+    "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"
+}

helpers/processor.py

@@ -1,32 +1,26 @@
+import io
 import cv2
 import imageio
 import numpy as np
 import torch
-import io
-
+from typing import Dict, List
+from fvcore.common.config import CfgNode
 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.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
-)
-from densepose.vis.extractor import (
-    CompoundExtractor,
-    create_extractor,
-    DensePoseResultExtractor
-)
+from densepose.vis.densepose_results_textures import DensePoseResultsVisualizerWithTexture as dp_iuv_texture
+from densepose.vis.extractor import CompoundExtractor, create_extractor, DensePoseResultExtractor
 
 class TextureProcessor:
-  def __init__(self, config, weights):
+  def __init__(self, config: str, weights: str) -> None:
     self.config = self.get_config(config, weights)
     self.predictor = DefaultPredictor(self.config)
     self.extractor = DensePoseResultExtractor()
   
-  def process_texture(self, image):
+  def process_texture(self, image: np.ndarray) -> np.ndarray:
     image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
     output = self.execute(image)
     if 'pred_densepose' in output:
@@ -44,7 +38,7 @@ class TextureProcessor:
     texture_atlas = self.process_texture(model_img)
     return self.overlay_texture(texture_atlas, person_img)
 
-  def overlay_texture(self, texture_atlas, original_image):
+  def overlay_texture(self, texture_atlas: np.ndarray, original_image: np.ndarray) -> np.ndarray:
     texture_atlas[:, :, :3] = texture_atlas[:, :, 2::-1]
     texture_atlases_dict = get_texture_atlases(None)
     vis = dp_iuv_texture(
@@ -68,17 +62,17 @@ class TextureProcessor:
 
     return image_vis
 
-  def parse_iuv(self, result):
+  def parse_iuv(self, result: Dict) -> np.ndarray:
     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):
+  def parse_bbox(self, result: Dict) -> np.ndarray:
     return result['pred_boxes_XYXY'][0].cpu().numpy()
 
-  def interpolate_tex(self, tex):
+  def interpolate_tex(self, tex: np.ndarray) -> np.ndarray:
     valid_mask = np.array((tex.sum(0) != 0) * 1, dtype='uint8')
     radius_increase = 10
     kernel = np.ones((radius_increase, radius_increase), np.uint8)
@@ -96,12 +90,18 @@ class TextureProcessor:
     
     return actual_part
 
-  def concat_textures(self, array):
+  def concat_textures(self, array: List[np.ndarray]) -> np.ndarray:
     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):
+  def get_texture(
+        self, 
+        im: np.ndarray, 
+        iuv: np.ndarray, 
+        bbox: List[int], 
+        tex_part_size: int = 200) -> np.ndarray:
+    
     im = im.transpose(2, 1, 0) / 255
     image_w, image_h = im.shape[1], im.shape[2]
     bbox[2] = bbox[2] - bbox[0]
@@ -142,14 +142,14 @@ class TextureProcessor:
 
     return tex
 
-  def create_iuv(self, results, image):
+  def create_iuv(self, results: Dict, image: np.ndarray) -> np.ndarray:
     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):
+  def get_config(self, config_fpath: str, model_fpath: str) -> CfgNode:
     cfg = get_cfg()
     add_densepose_config(cfg)
     cfg.merge_from_file(config_fpath)
@@ -158,12 +158,12 @@ class TextureProcessor:
     cfg.freeze()
     return cfg
 
-  def execute(self, image):
+  def execute(self, image: np.ndarray) -> Dict:
     with torch.no_grad():
         outputs = self.predictor(image)['instances']
         return self.execute_on_outputs(outputs)
 
-  def execute_on_outputs(self, outputs: Instances):
+  def execute_on_outputs(self, outputs: Instances) -> Dict:
     result = {}
     if outputs.has('scores'):
         result['scores'] = outputs.get('scores').cpu()