Update app.py

app.py

@@ -79,7 +79,7 @@ def detect(img,thr=0.2,trained_dataset='aitod'):
   t_img = transformer(img).unsqueeze(0)#.unsqueeze(0)      #[1,1,640,640]
   size = torch.tensor([[t_img.shape[2], t_img.shape[3]]])
   #print(t_img.shape)
-  if cope == 'aitod':
+  if trained_dataset == 'aitod':
       labels, boxes, scores=model(t_img,size)
   elif cope == 'cope12':
       labels, boxes, scores=model2(t_img,size)
@@ -110,7 +110,7 @@ def detect(img,thr=0.2,trained_dataset='aitod'):
            90: 'toothbrush'}                #coco
 
       label_color_dict = {8:'burlyWood',7:'red',6:'blue',5:'green',4:'yellow',3:'cyan',2:'magenta',1:'orange'}
-      if cope == 'coco':
+      if trained_dataset == 'coco':
           for idx,b in enumerate(box):
               label_i = lab[idx].item()
               draw.rectangle(list(b), outline=label_color_dict[label_i], )
@@ -124,6 +124,6 @@ def detect(img,thr=0.2,trained_dataset='aitod'):
   #save_path = Path('./output') / img_path.name
   return img
 
-interface = gr.Interface(fn=detect,inputs=["image",gr.Slider(label="thr", value=0.2, maximum=1, minimum=0),gr.Radio(['aitod','ten_classes','COCO'])],outputs="image",title="rt-cope detect")
+interface = gr.Interface(fn=detect,inputs=["image",gr.Slider(label="thr", value=0.2, maximum=1, minimum=0),gr.Radio(['aitod','ten_classes','COCO'],value='aitod')],outputs="image",title="rt-cope detect")
   
 interface.launch()