multi-categories

app.py

@@ -4,12 +4,14 @@ import torch
 from torchvision.transforms import InterpolationMode
 
 BICUBIC = InterpolationMode.BICUBIC
-from utils import setup, get_similarity_map, display_segmented_sketch
+from utils import setup, get_similarity_map, display_segmented_sketch,get_noun_phrase
 from vpt.launch import default_argument_parser
 from collections import OrderedDict
 import numpy as np
 import matplotlib.pyplot as plt
 import models
+import string
+import nltk
 import torchvision
 
 args = default_argument_parser().parse_args()
@@ -31,10 +33,18 @@ print("Model loaded successfully")
 
 def run(sketch, caption, threshold, seed):
     # set the condidate classes here
-    classes = [caption]
-
-    colors = plt.get_cmap("tab10").colors
-    classes_colors = colors[3:len(classes) + 3]
+    caption = caption.replace('\n',' ')
+    translator = str.maketrans('', '', string.punctuation)
+    caption = caption.translate(translator).lower()
+    words = nltk.word_tokenize(caption)
+    classes = get_noun_phrase(words)
+    if len(classes) ==0:
+        classes = [caption]
+
+    # print(classes)
+    
+    colors = plt.get_cmap("Set1").colors
+    classes_colors = colors[:len(classes)]
 
     sketch2 = sketch['composite']   
 
@@ -414,7 +424,7 @@ with gr.Blocks(js=scripts, css=css, theme='gstaff/xkcd') as demo:
         examples_per_page=30,
         examples=[
         ['demo/sketch_1.png', 'giraffe looking at you', 0.6],
-        ['demo/sketch_2.png', 'tree on the right', 0.6],
+        ['demo/sketch_2.png', 'a kite flying in the sky', 0.6],
         ['demo/sketch_3.png', 'a girl playing', 0.6],
         ['demo/000000004068.png', 'car going so fast', 0.6],
         ['demo/000000004546.png', 'mountains in the background', 0.6],
@@ -429,7 +439,7 @@ with gr.Blocks(js=scripts, css=css, theme='gstaff/xkcd') as demo:
         ['demo/000000305414.png', 'a lonely elephant roaming around', 0.6],
         ['demo/000000484246.png', 'giraffe with a loong neck', 0.6],
         ['demo/000000549338.png', 'two donkeys trying to be smart', 0.6],
-        ['demo/000000038116.png', 'a bat on the left', 0.6],
+        ['demo/000000038116.png', 'a bat next to a kid', 0.6],
         ['demo/000000221509.png', 'funny looking cow', 0.6],
         ['demo/000000246066.png', 'bench in the park', 0.6],
         ['demo/000000001611.png', 'trees in the background', 0.6]
@@ -439,8 +449,4 @@ with gr.Blocks(js=scripts, css=css, theme='gstaff/xkcd') as demo:
         # cache_examples=True,
     )
 
-
-
-
-
-demo.launch(share=False, )
+demo.launch(share=False)

utils.py

@@ -10,11 +10,41 @@ from vpt.src.utils.file_io import PathManager
 import matplotlib.pyplot as plt
 from matplotlib.colors import rgb_to_hsv, hsv_to_rgb
 import warnings
+import nltk
 
 
 warnings.filterwarnings("ignore")
 
 
+def get_noun_phrase(tokenized):
+    # Taken from Su Nam Kim Paper...
+    grammar = r"""
+        NBAR:
+            {<NN.*|JJ>*<NN.*>}  # Nouns and Adjectives, terminated with Nouns
+        NP:
+            {<NBAR>}
+            {<NBAR><IN><NBAR>}  # Above, connected with in/of/etc...
+    """
+    chunker = nltk.RegexpParser(grammar)
+
+    chunked = chunker.parse(nltk.pos_tag(tokenized))
+    continuous_chunk = []
+    current_chunk = []
+
+    for subtree in chunked:
+        if isinstance(subtree, nltk.Tree):
+            current_chunk.append(' '.join([token for token, pos in subtree.leaves()]))
+        elif current_chunk:
+            named_entity = ' '.join(current_chunk)
+            if named_entity not in continuous_chunk:
+                continuous_chunk.append(named_entity)
+                current_chunk = []
+        else:
+            continue
+
+    return continuous_chunk
+
+
 def setup(args):
     """
     Create configs and perform basic setups.
@@ -90,14 +120,16 @@ def display_segmented_sketch(pixel_similarity_array,binary_sketch,classes,classe
     # Convert the HSV image back to RGB to display and save
     rgb_image = hsv_to_rgb(hsv_image)
 
-    # # Calculate centroids and render class names
-    # for i, class_name in enumerate(classes):
-    #     mask = class_indices == i
-    #     if np.any(mask):
-    #         y, x = np.nonzero(mask)
-    #         centroid_x, centroid_y = np.mean(x), np.mean(y)
-    #         plt.text(centroid_x, centroid_y, class_name, color=classes_colors[i], ha='center', va='center',fontsize=14,   # color=classes_colors[i]
-    #         bbox=dict(facecolor='lightgrey', edgecolor='none', boxstyle='round,pad=0.2', alpha=0.8))
+    
+    if len(classes) > 1:
+        # Calculate centroids and render class names
+        for i, class_name in enumerate(classes):
+            mask = class_indices == i
+            if np.any(mask):
+                y, x = np.nonzero(mask)
+                centroid_x, centroid_y = np.mean(x), np.mean(y)
+                plt.text(centroid_x, centroid_y, class_name, color=classes_colors[i], ha='center', va='center',fontsize=10,   # color=classes_colors[i]
+                bbox=dict(facecolor='lightgrey', edgecolor='none', boxstyle='round,pad=0.2', alpha=0.8))
 
 
     # Display the image with class names