multi_classes is False

app.py

@@ -20,6 +20,8 @@ import torchvision
 args = default_argument_parser().parse_args()
 cfg = setup(args)
 
+multi_classes = False
+
 device = "cpu"  # "cuda" if torch.cuda.is_available() else "cpu"
 Ours, preprocess = models.load("CS-ViT-B/16", device=device, cfg=cfg, train_bool=False)
 state_dict = torch.load("sketch_seg_best_miou.pth", map_location=device)
@@ -41,7 +43,7 @@ def run(sketch, caption, threshold, seed):
     caption = caption.translate(translator).lower()
     words = word_tokenize(caption)
     classes = get_noun_phrase(words)
-    if len(classes) ==0:
+    if len(classes) ==0 or multi_classes == False:
         classes = [caption]
 
     # print(classes)
@@ -81,6 +83,7 @@ def run(sketch, caption, threshold, seed):
     pixel_similarity[pixel_similarity < threshold] = 0
     pixel_similarity_array = pixel_similarity.cpu().numpy().transpose(2, 0, 1)
 
+
     # display_segmented_sketch(pixel_similarity_array, sketch2, classes, classes_colors, live=True)
     
     # Find the class index with the highest similarity for each pixel