Browse Source

release YOLOF-R18

yjh0410 1 year ago
parent
commit
417bebcb30
2 changed files with 9 additions and 3 deletions
  1. 4 2
      odlab/evaluator/coco_evaluator.py
  2. 5 1
      odlab/models/detectors/fcos/fcos.py

+ 4 - 2
odlab/evaluator/coco_evaluator.py

@@ -39,7 +39,9 @@ class COCOAPIEvaluator():
             # inference
             image = image.unsqueeze(0).to(self.device)
             outputs = model(image)
-            bboxes, scores, cls_inds = outputs
+            scores = outputs['scores']
+            labels = outputs['labels']
+            bboxes = outputs['bboxes']
 
             # rescale bbox
             orig_h, orig_w = target["orig_size"].tolist()
@@ -52,7 +54,7 @@ class COCOAPIEvaluator():
                 y1 = float(box[1])
                 x2 = float(box[2])
                 y2 = float(box[3])
-                label = self.dataset.coco_indexs[int(cls_inds[i])]
+                label = self.dataset.coco_indexs[int(labels[i])]
                 
                 # COCO json format
                 bbox = [x1, y1, x2 - x1, y2 - y1]

+ 5 - 1
odlab/models/detectors/fcos/fcos.py

@@ -117,6 +117,10 @@ class FCOS(nn.Module):
             bboxes[..., 1::2] /= src.shape[-2]
             bboxes = bboxes.clip(0., 1.)
 
-            return bboxes, scores, labels
+            outputs = {
+                'scores': scores,
+                'labels': labels,
+                'bboxes': bboxes
+            }
 
         return outputs