Update app.py

app.py

@@ -41,10 +41,9 @@ transformer = transforms.Compose([
 class Model(nn.Module):
     def __init__(self, confg=None, ckpt="",cope=True) -> None:
         super().__init__()
-        if cope:
-            self.cfg = src.core.YAMLConfig(confg, resume=ckpt)
-        else:
-            self.cfg = src2.core.YAMLConfig(confg, resume=ckpt)
+
+        self.cfg = src.core.YAMLConfig(confg, resume=ckpt)
+
         if ckpt:
             checkpoint = torch.load(ckpt, map_location='cpu') 
             if 'ema' in checkpoint:
@@ -66,8 +65,9 @@ class Model(nn.Module):
         return self.postprocessor(outputs, orig_target_sizes)
   
 
-model = Model(confg='./configs/rtdetr/rtdetr_r101vd_6x_coco.yml',ckpt="./checkpointcope12.pth")
-model2 = Model(confg='./configs/rtdetr/rtdetr_r101vd_6x_coco.yml',ckpt="./checkpoint_init.pth",cope=False)
+model = Model(confg='./configs/rtdetr/rtdetr_r101vd_6x_coco.yml',ckpt="./checkpoint_init.pth")
+model2 = Model(confg='./configs/rtdetr/rtdetr_r101vd_6x_coco_cope12.yml',ckpt="./checkpointcope12.pth",cope=False)
+model3 = Model(confg='./configs/rtdetr/rtdetr_r101vd_6x_coco_cope24.yml',ckpt="./checkpointcope24.pth",cope=False)
 
 #img = cv2.imread('./j.jpg',cv2.IMREAD_GRAYSCALE)
 #img = Image.open('./a.jpg').convert('RGB').resize((640,640))
@@ -79,8 +79,12 @@ def detect(img,thr=0.2,cope='none'):
   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)
-  labels, boxes, scores=model(t_img,size)
-  labels2, boxes2, scores2=model2(t_img,size)
+  if cope == 'none':
+      labels, boxes, scores=model(t_img,size)
+  elif cope == 'cope12':
+      labels, boxes, scores=model2(t_img,size)
+  else
+      labels, boxes, scores=model3(t_img,size)
   #img_path = Path('./a.jpg')
   draw = ImageDraw.Draw(img)
   thrh = thr
@@ -101,6 +105,6 @@ def detect(img,thr=0.2,cope='none'):
   #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)],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.inputs.Radio(['none','cope12','cope24'])],outputs="image",title="rt-cope detect")
   
 interface.launch()