optimize

app.py

@@ -82,10 +82,7 @@ def save_masks(preds, text):
         cv2.imwrite(dir, mask * 255)
 
 def predict(image, text):
-    args = get_parser().parse_args()
-    cfg = setup_cfg(args)
-    demo = VisualizationDemo(cfg, text=text)
-    predictions, visualized_output = demo.run_on_image(image)
+    predictions, visualized_output = demo.run_on_image(image, text)
     #save_masks(predictions, text.split(','))
     canvas = fc(visualized_output.fig)
     canvas.draw()
@@ -96,6 +93,8 @@ def predict(image, text):
 if __name__ == "__main__":
     args = get_parser().parse_args()
     cfg = setup_cfg(args)
+    global demo
+    demo = VisualizationDemo(cfg)
 
     iface = gr.Interface(
         fn=predict,

demo/predictor.py

@@ -49,7 +49,7 @@ class VisualizationDemo(object):
             self.metadata = ns()
             self.metadata.stuff_classes = pred.test_class_texts
 
-    def run_on_image(self, image):
+    def run_on_image(self, image, text=None):
         """
         Args:
             image (np.ndarray): an image of shape (H, W, C) (in BGR order).
@@ -59,6 +59,17 @@ class VisualizationDemo(object):
             vis_output (VisImage): the visualized image output.
         """
         vis_output = None
+
+        if text is not None:
+            pred = self.predictor.model.sem_seg_head.predictor
+            pred.test_class_texts = text.split(',')
+            pred.text_features_test = pred.class_embeddings(pred.test_class_texts, 
+                #imagenet_templates.IMAGENET_TEMPLATES, 
+                 ['A photo of a {} in the scene',],
+                pred.clip_model).permute(1, 0, 2).float().repeat(1, 80, 1)
+            self.metadata = ns()
+            self.metadata.stuff_classes = pred.test_class_texts
+
         predictions = self.predictor(image)
         # Convert image from OpenCV BGR format to Matplotlib RGB format.
         image = image[:, :, ::-1]