Update app.py

support multitask

app.py

@@ -11,7 +11,7 @@ import torch
 import numpy as np
 import sys
 import os
-version = "YxZhang/evf-sam"
+version = "YxZhang/evf-sam-multitask"
 model_type = "ori"
 
 tokenizer = AutoTokenizer.from_pretrained(
@@ -29,7 +29,7 @@ model.to('cuda')
 
 @spaces.GPU
 @torch.no_grad()
-def pred(image_np, prompt):
+def pred(image_np, prompt, semantic_type):
     original_size_list = [image_np.shape[:2]]
 
     image_beit = beit3_preprocess(image_np, 224).to(dtype=model.dtype,
@@ -38,6 +38,8 @@ def pred(image_np, prompt):
     image_sam, resize_shape = sam_preprocess(image_np, model_type=model_type)
     image_sam = image_sam.to(dtype=model.dtype, device=model.device)
 
+    if semantic_type:
+        prompt = "[semantic] " + prompt
     input_ids = tokenizer(
         prompt, return_tensors="pt")["input_ids"].to(device=model.device)
 
@@ -79,7 +81,12 @@ demo = gr.Interface(
             label="Prompt",
             info=
             "Use a phrase or sentence to describe the object you want to segment. Currently we only support English"
-        )
+        ),
+        gr.components.Checkbox(
+            False, 
+            label="semantic level", 
+            info="check this if you want to segment body parts or background or multi objects (only available with latest evf-sam checkpoint)"
+        ),
     ],
     outputs=[
         gr.components.Image(type="numpy", label="visulization"),