add multi-class

app.py

@@ -17,10 +17,31 @@ nltk.download('averaged_perceptron_tagger')
 from nltk.tokenize import word_tokenize
 import torchvision
 
+import spacy
+
+# download the model
+spacy.cli.download("en_core_web_sm")
+
+# Load spaCy model
+nlp = spacy.load("en_core_web_sm")
+
+def extract_objects(prompt):
+    doc = nlp(prompt)
+    # Extract object nouns (including proper nouns and compound nouns)
+    objects = set()
+    for token in doc:
+        # Check if the token is a noun or part of a named entity
+        if token.pos_ in {"NOUN", "PROPN"} or token.ent_type_:
+            objects.add(token.text)
+        # Check if the token is part of a compound noun
+        if token.dep_ in {"compound"}:
+            objects.add(token.head.text)
+    return list(objects)
+
 args = default_argument_parser().parse_args()
 cfg = setup(args)
 
-multi_classes = False
+multi_classes = True
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 Ours, preprocess = models.load("CS-ViT-B/16", device=device, cfg=cfg, train_bool=False)
@@ -42,10 +63,12 @@ def run(sketch, caption, threshold, seed):
     
     # set the condidate classes here
     caption = caption.replace('\n',' ')
-    translator = str.maketrans('', '', string.punctuation)
-    caption = caption.translate(translator).lower()
-    words = word_tokenize(caption)
-    classes = get_noun_phrase(words)
+    classes = extract_objects(caption)
+    # translator = str.maketrans('', '', string.punctuation)
+    # caption = caption.translate(translator).lower()
+    # words = word_tokenize(caption)
+    # classes = get_noun_phrase(words)
+    # print(classes)
     if len(classes) ==0 or multi_classes == False:
         classes = [caption]
 

requirements.txt

@@ -10,4 +10,5 @@ iopath
 ftfy
 fvcore
 regex
-nltk
+nltk
+spacy