From 699fa88e40fdbc819631d23ff1c5ea33745b11a2 Mon Sep 17 00:00:00 2001
From: smthemex <138738845+smthemex@users.noreply.github.com>
Date: Wed, 12 Jun 2024 19:45:37 +0800
Subject: [PATCH] Add files via upload

add
---
 utils/clothes_utils.py | 26 ++++++++++++++++++
 utils/vis_utils.py     | 60 ++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 86 insertions(+)
 create mode 100644 utils/clothes_utils.py
 create mode 100644 utils/vis_utils.py

diff --git a/utils/clothes_utils.py b/utils/clothes_utils.py
new file mode 100644
index 0000000..4ddcdeb
--- /dev/null
+++ b/utils/clothes_utils.py
@@ -0,0 +1,26 @@
+ind2labels_dict = {
+  0: "Background",
+  1: "Hat",
+  2: "Hair",
+  3: "Sunglasses",
+  4: "Upper-clothes",
+  5: "Skirt",
+  6: "Pants",
+  7: "Dress",
+  8: "Belt",
+  9: "Left-shoe",
+  10: "Right-shoe",
+  11: "Face",
+  12: "Left-leg",
+  13: "Right-leg",
+  14: "Left-arm",
+  15: "Right-arm",
+  16: "Bag",
+  17: "Scarf"
+}
+
+# Calculate the reverse dict
+labels2ind_dict = {v: k for k, v in ind2labels_dict.items()}
+
+
+relevant_inds = [1,2,3,4,5,6,7,8,9,10,16,17]
diff --git a/utils/vis_utils.py b/utils/vis_utils.py
new file mode 100644
index 0000000..62f2e0b
--- /dev/null
+++ b/utils/vis_utils.py
@@ -0,0 +1,60 @@
+import textwrap
+from typing import List, Tuple, Optional
+
+import numpy as np
+from PIL import Image, ImageDraw, ImageFont
+
+LINE_WIDTH = 20
+
+
+def add_text_to_image(image: np.ndarray, text: str, text_color: Tuple[int, int, int] = (0, 0, 0),
+                      min_lines: Optional[int] = None, add_below: bool = True):
+    import textwrap
+    lines = textwrap.wrap(text, width=LINE_WIDTH)
+    if min_lines is not None and len(lines) < min_lines:
+        if add_below:
+            lines += [''] * (min_lines - len(lines))
+        else:
+            lines = [''] * (min_lines - len(lines)) + lines
+    h, w, c = image.shape
+    offset = int(h * .12)
+    img = np.ones((h + offset * len(lines), w, c), dtype=np.uint8) * 255
+    font_size = int(offset * .8)
+
+    try:
+        font = ImageFont.truetype("assets/OpenSans-Regular.ttf", font_size)
+        textsize = font.getbbox(text)
+        y_offset = (offset - textsize[3]) // 2
+    except:
+        font = ImageFont.load_default()
+        y_offset = offset // 2
+
+    if add_below:
+        img[:h] = image
+    else:
+        img[-h:] = image
+    img = Image.fromarray(img)
+    draw = ImageDraw.Draw(img)
+    for i, line in enumerate(lines):
+        line_size = font.getbbox(line)
+        text_x = (w - line_size[2]) // 2
+        if add_below:
+            draw.text((text_x, h + y_offset + offset * i), line, font=font, fill=text_color)
+        else:
+            draw.text((text_x, 0 + y_offset + offset * i), line, font=font, fill=text_color)
+    return np.array(img)
+
+
+def create_table_plot(images: List[Image.Image], titles: List[str]=None, captions: List[str]=None) -> Image.Image:
+    title_max_lines = np.max([len(textwrap.wrap(text, width=LINE_WIDTH)) for text in titles]) if titles is not None else 0
+    caption_max_lines = np.max([len(textwrap.wrap(text, width=LINE_WIDTH)) for text in captions]) if captions is not None else 0
+    out_images = []
+    for i in range(len(images)):
+        im = np.array(images[i])
+        if titles is not None:
+            im = add_text_to_image(im, titles[i], add_below=False, min_lines=title_max_lines)
+        if captions is not None:
+            im = add_text_to_image(im, captions[i], add_below=True, min_lines=caption_max_lines)
+        out_images.append(im)
+    image = Image.fromarray(np.concatenate(out_images, axis=1))
+    return image