Update app.py

app.py

@@ -1,15 +1,30 @@
 '''
 contains the infernce code for fastsam.py
 '''
+import os 
+import sys 
+import cv2 
+import torch
 import numpy as np 
 import gradio as gr 
 from PIL import Image 
+import matplotlib.pyplot as plt 
 from ultralytics import FastSAM 
 from ultralytics.models.fastsam import FastSAMPrompt 
 
 def inference(image):
 
-    model = FastSAM('FastSAM.pt')
+    #model = FastSAM('FastSAM.pt')
+    model = FastSAM(params['./FastSAM.pt'])
+    device = torch.device("cuda") if torch.cuda.is_available() else "cpu"
+    text_prompt = params['text_prompt']
+    input_image = Image.fromarray(input_image).convert("RGB")
+    w,h = input_image.size
+    scale = params['input_size']/max(w,h) 
+
+    new_w,new_h = int(w*scale),int(h*scale) 
+    input_image = input_image.resize((new_w,new_h))
+    
 
     ## inference results
     # Run inference on an image
@@ -17,29 +32,33 @@ def inference(image):
 
     # Prepare a Prompt Process object
     prompt_process = FastSAMPrompt(source, results, device='cpu')
-    
-    # Everything prompt
-    ann = prompt_process.everything_prompt()
 
-     # Bbox default shape [0,0,0,0] -> [x1,y1,x2,y2]
-    ann = prompt_process.box_prompt(bbox=[200, 200, 300, 300])
-    
-    # Text prompt
-    ann = prompt_process.text_prompt(text='a photo of a dog')
-    
-    # Bbox default shape [0,0,0,0] -> [x1,y1,x2,y2]
-    ann = prompt_process.point_prompt(points=[[200,200]],pointlabel=[1])
-    return prompt_process.plot(annotations=ann,output='./')
-    
+    if text_prompt:
+        ann = prompt_process.text_prompt(text=text_prompt)
+    if show_all:
+        gr.Warning("Annotations are enabled, Overiding the text prompt!")
+        ann = prompt_process.everything_prompt()
 
-    
+    result = prompt_process.plot_to_result(annotations=ann,bboxes=None,points=None,point_label=None,withContours=True,better_quality=False,)
+
+    return result 
+
+
+params = {
+            'input_size':1024,
+            'point_prompt':[[0,0]],
+            'text_prompt':'jeep',
+            'output_dir':'./testimg.png'
+}
 
 title = "Usage of FastSAM"
 description = "Implementation of pre-trained fast-sam model for spaces."
 
-demo = gr.Interface(inference,inputs=[gr.Image(shape=(32,32),labels="Input Image"),"checkbox"],
+demo = gr.Interface(inference,inputs=[gr.Image(shape=(32,32),labels="Input Image"),
+                                      gr.Checkbox(True,label='Show All Annotations'),
+                                      gr.Textbox('',label='Provide the text prompt: '),],
                     outputs = [gr.Image(shape=(32,32),label='Output').style(width=128,height=128)],
-                    title = title, 
+                    title = title,
                     description = description)
 
 demo.launch(debug = True)