Adding text promp

app.py

@@ -14,33 +14,41 @@ elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
 model = FastSAM('./weights/FastSAM-x.pt')
 model.to(device)
 
-def inference(image, conf_thres, iou_thres,):
+def inference(image, conf_thres, iou_thres, text):
     pred = model(image, device=device, retina_masks=True, imgsz=1024, conf=conf_thres, iou=iou_thres)
     prompt_process = FastSAMPrompt(image, pred, device="cpu")
     ann = prompt_process.everything_prompt()
-    prompt_process.plot(annotations=ann, output_path="./output.jpg", withContours=False, better_quality=False)
-    output = Image.open('./output.jpg')
-    output = np.array(output)
-    return output
+    prompt_process.plot(annotations=ann, output_path="./output_sam.jpg", withContours=False, better_quality=False)
+    output_sam = Image.open('./output_sam.jpg')
+    output_sam = np.array(output_sam)
+    output_text = None
+    if text:
+        ann = prompt_process.text_prompt(text=text)
+        prompt_process.plot(annotations=ann, output_path="./output_text.jpg", withContours=False, better_quality=False)
+        output_text = Image.open('./output_text.jpg')
+        output_text = np.array(output_text)
+    return output_sam, output_text
 
 title = "FAST-SAM Segment Anything"
 description = "A simple Gradio interface to infer on FAST-SAM model"
-examples = [["image_1.jpg", 0.25, 0.45],
-            ["image_2.jpg", 0.25, 0.45],
-            ["image_3.jpg", 0.25, 0.45],
-            ["image_4.jpg", 0.25, 0.45],
-            ["image_5.jpg", 0.25, 0.45],
-            ["image_6.jpg", 0.25, 0.45],
-            ["image_7.jpg", 0.25, 0.45],
-            ["image_8.jpg", 0.25, 0.45],
-            ["image_9.jpg", 0.25, 0.45],
-            ["image_10.jpg", 0.25, 0.45]]
+examples = [["image_1.jpg", 0.25, 0.45, 'A black tire'],
+            ["image_2.jpg", 0.25, 0.45, 'Shades of blue'],
+            ["image_3.jpg", 0.25, 0.45, 'A spiral staircase'],
+            ["image_4.jpg", 0.25, 0.45, 'A clock and a plane'],
+            ["image_5.jpg", 0.25, 0.45, 'Clouds in the sky'],
+            ["image_6.jpg", 0.25, 0.45, 'Front wheel'],
+            ["image_7.jpg", 0.25, 0.45, 'A white chair'],
+            ["image_8.jpg", 0.25, 0.45, 'The grassy field'],
+            ["image_9.jpg", 0.25, 0.45, 'Rock formation'],
+            ["image_10.jpg", 0.25, 0.45, 'A rope railing']]
 
 demo = gr.Interface(inference, 
-                    inputs = [gr.Image(width=320, height=320, label="Input Image"), 
+                    inputs = [gr.Image(width=640, height=640, label="Input Image"), 
                               gr.Slider(0, 1, 0.25, label="Confidence Threshold"),
-                              gr.Slider(0, 1, 0.45, label="IoU Thresold")],
-                    outputs= [gr.Image(width=640, height=640, label="Output")],
+                              gr.Slider(0, 1, 0.45, label="IoU Thresold"),
+                              gr.Textbox(label="Enter text promp", type="text"),],
+                    outputs= [gr.Image(width=640, height=640, label="Output SAM"),
+                              gr.Image(width=640, height=640, label="Output text prompt")],
                     title=title,
                     description=description,
                     examples=examples)