Update app.py

app.py

@@ -34,7 +34,6 @@ from data import get_dataset
 import torchvision.transforms as transforms
 
 import gradio as gr
-import streamlit as st
 
 model_name = "convnext_xlarge_in22k"
 device = 'cuda' if torch.cuda.is_available() else 'cpu'
@@ -332,57 +331,42 @@ def load_model():
 """
 lseg_model, lseg_transform = load_model()
 
-# to be revised
-uploaded_file = gr.inputs.Image(type='pil')
-input_labels = st.text_input("Input labels", value="dog, grass, other")
-gr.outputs.Label(type="confidences",num_top_classes=5)
-st.write("The labels are", input_labels)
+def inference(image,text):
+    input_labels = text
 
-image = Image.open(uploaded_file)
-pimage = lseg_transform(np.array(image)).unsqueeze(0)
+    pimage = lseg_transform(np.array(image)).unsqueeze(0)
 
-labels = []
-for label in input_labels.split(","):
-    labels.append(label.strip())
-
-with torch.no_grad():
-    outputs = lseg_model.parallel_forward(pimage, labels)
+    labels = []
+    for label in input_labels.split(","):
+        labels.append(label.strip())
     
-    predicts = [
-        torch.max(output, 1)[1].cpu().numpy()
-        for output in outputs
-    ]
+    with torch.no_grad():
+        outputs = lseg_model.parallel_forward(pimage, labels)
+        
+        predicts = [
+            torch.max(output, 1)[1].cpu().numpy()
+            for output in outputs
+        ]
+        
+    image = pimage[0].permute(1,2,0)
+    image = image * 0.5 + 0.5
+    image = Image.fromarray(np.uint8(255*image)).convert("RGBA")
     
-image = pimage[0].permute(1,2,0)
-image = image * 0.5 + 0.5
-image = Image.fromarray(np.uint8(255*image)).convert("RGBA")
-
-pred = predicts[0]
-new_palette = get_new_pallete(len(labels))
-mask, patches = get_new_mask_pallete(pred, new_palette, out_label_flag=True, labels=labels)
-seg = mask.convert("RGBA")
-
-fig = plt.figure()
-plt.subplot(121)
-plt.imshow(image)
-plt.axis('off')
-
-plt.subplot(122)
-plt.imshow(seg)
-plt.legend(handles=patches, loc='upper right', bbox_to_anchor=(1.3, 1), prop={'size': 5})
-plt.axis('off')
-
-plt.tight_layout()
-
-#st.image([image,seg], width=700, caption=["Input image", "Segmentation"])
-st.pyplot(fig)
-
-title = "LSeg"
-
-description = "Gradio demo for LSeg for semantic segmentation. To use it, simply upload your image, or click one of the examples to load them, then add any label set"
-
-article = "<p style='text-align: center'><a href='https://arxiv.org/abs/2201.03546' target='_blank'>Language-driven Semantic Segmentation</a> | <a href='hhttps://github.com/isl-org/lang-seg' target='_blank'>Github Repo</a></p>"
+    pred = predicts[0]
+    new_palette = get_new_pallete(len(labels))
+    mask, patches = get_new_mask_pallete(pred, new_palette, out_label_flag=True, labels=labels)
+    seg = mask.convert("RGBA")
+
+    fig = plt.figure()
+    plt.subplot(121)
+    plt.axis('off')
+
+    plt.subplot(122)
+    plt.legend(handles=patches, loc='upper right', bbox_to_anchor=(1.3, 1), prop={'size': 5})
+    plt.axis('off')
+    
+    plt.tight_layout()
 
-examples = ['test.jpeg']
+    return plt    
 
-gr.Interface(inference, inputs, outputs, title=title, description=description, article=article, analytics_enabled=False, examples=examples).launch(enable_queue=True)
+gr.Interface(inference,["image","text"],"plot").launch(debug=True)