Update app.py

app.py

@@ -16,11 +16,18 @@ from utils.torch_utils import select_device, load_classifier, time_synchronized,
 
 
 os.system('git clone https://github.com/WongKinYiu/yolov7')
-
-
- 
-def Custom_detect(img):
-    model='best' 
+os.system('wget https://github.com/WongKinYiu/yolov7/releases/download/v0.1/yolov7.pt')
+os.system('wget https://github.com/WongKinYiu/yolov7/releases/download/v0.1/yolov7-mask.pt')
+
+#model='best'  
+def Custom_detect(img,mode):
+    if mode=='Custom Detection':
+        model='best'
+    elif mode=='Instance Segmentation':
+        model='yolov7-mask'
+    else:
+        model='yolov7'
+    
     parser = argparse.ArgumentParser()
     parser.add_argument('--weights', nargs='+', type=str, default=model+".pt", help='model.pt path(s)')
     parser.add_argument('--source', type=str, default='Temp_files/', help='source') 
@@ -153,10 +160,11 @@ def Custom_detect(img):
 
     return Image.fromarray(im0[:,:,::-1])
 inp = gr.Image(type="pil")
+inp2= gr.Dropdown((default="Custom Detection",choices=["Custom Detection","Instance Segmentation","Yolov7 model detection"])]
 output = gr.Image(type="pil")
 
 examples=["Examples/Image1.jpg","Examples/Image14.jpg","Examples/Image32.jpg"]
 
-io=gr.Interface(fn=Custom_detect, inputs=inp, outputs=output, title='Vehicle Detection With Custom YOLOv7',examples=examples,cache_examples=False)
+io=gr.Interface(fn=Custom_detect, inputs=[inp,inp2], outputs=output, title='Vehicle Detection With Custom YOLOv7',examples=examples,cache_examples=False)
 io.launch()