Update app.py

app.py

@@ -81,7 +81,7 @@ def detect(img,thr=0.2,trained_dataset='aitod'):
   #print(t_img.shape)
   if trained_dataset == 'aitod':
       labels, boxes, scores=model(t_img,size)
-  elif trained_dataset == 'cope12':
+  elif trained_dataset == 'ten_classes':
       labels, boxes, scores=model2(t_img,size)
   else:
       labels, boxes, scores=model3(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 trained_dataset == '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], )