Forráskód Böngészése

add fake batch prediction for rcnn models

will-jl944 4 éve
szülő
commit
a1ca94c3a8
3 módosított fájl, 70 hozzáadás és 21 törlés
  1. 33 8
      paddlex/cv/models/detector.py
  2. 12 2
      paddlex/cv/transforms/batch_operators.py
  3. 25 11
      paddlex/deploy.py

+ 33 - 8
paddlex/cv/models/detector.py

@@ -105,8 +105,6 @@ class BaseDetector(BaseModel):
         if mode in ['train', 'eval']:
             outputs = net_out
         else:
-            for key in ['im_shape', 'scale_factor']:
-                net_out[key] = inputs[key]
             outputs = dict()
             for key in net_out:
                 outputs[key] = net_out[key].numpy()
@@ -519,8 +517,14 @@ class BaseDetector(BaseModel):
         batch_transforms = self._compose_batch_transform(transforms, 'test')
         batch_samples = batch_transforms(batch_samples)
         if to_tensor:
-            for k, v in batch_samples.items():
-                batch_samples[k] = paddle.to_tensor(v)
+            if isinstance(batch_samples, dict):
+                for k in batch_samples:
+                    batch_samples[k] = paddle.to_tensor(batch_samples[k])
+            else:
+                for sample in batch_samples:
+                    for k in sample:
+                        sample[k] = paddle.to_tensor(sample[k])
+
         return batch_samples
 
     def _postprocess(self, batch_pred):
@@ -989,6 +993,18 @@ class FasterRCNN(BaseDetector):
         super(FasterRCNN, self).__init__(
             model_name='FasterRCNN', num_classes=num_classes, **params)
 
+    def run(self, net, inputs, mode):
+        if mode in ['train', 'eval']:
+            outputs = net(inputs)
+        else:
+            outputs = []
+            for sample in inputs:
+                net_out = net(sample)
+                for key in net_out:
+                    net_out[key] = net_out[key].numpy()
+                outputs.append(net_out)
+        return outputs
+
     def _compose_batch_transform(self, transforms, mode='train'):
         if mode == 'train':
             default_batch_transforms = [
@@ -1012,7 +1028,8 @@ class FasterRCNN(BaseDetector):
 
         batch_transforms = BatchCompose(
             custom_batch_transforms + default_batch_transforms,
-            collate_batch=collate_batch)
+            collate_batch=collate_batch,
+            return_list=mode == 'test')
 
         return batch_transforms
 
@@ -1058,6 +1075,13 @@ class FasterRCNN(BaseDetector):
         self.fixed_input_shape = image_shape
         return self._define_input_spec(image_shape)
 
+    def _postprocess(self, batch_pred):
+        prediction = [
+            super(FasterRCNN, self)._postprocess(pred)[0]
+            for pred in batch_pred
+        ]
+        return prediction
+
 
 class PPYOLO(YOLOv3):
     def __init__(self,
@@ -1479,7 +1503,7 @@ class PPYOLOv2(YOLOv3):
         return self._define_input_spec(image_shape)
 
 
-class MaskRCNN(BaseDetector):
+class MaskRCNN(FasterRCNN):
     def __init__(self,
                  num_classes=80,
                  backbone='ResNet50_vd',
@@ -1714,7 +1738,7 @@ class MaskRCNN(BaseDetector):
                 'mask_post_process': mask_post_process
             })
         self.with_fpn = with_fpn
-        super(MaskRCNN, self).__init__(
+        super(FasterRCNN, self).__init__(
             model_name='MaskRCNN', num_classes=num_classes, **params)
 
     def _compose_batch_transform(self, transforms, mode='train'):
@@ -1740,7 +1764,8 @@ class MaskRCNN(BaseDetector):
 
         batch_transforms = BatchCompose(
             custom_batch_transforms + default_batch_transforms,
-            collate_batch=collate_batch)
+            collate_batch=collate_batch,
+            return_list=mode == 'test')
 
         return batch_transforms
 

+ 12 - 2
paddlex/cv/transforms/batch_operators.py

@@ -26,10 +26,14 @@ from paddlex.utils import logging
 
 
 class BatchCompose(Transform):
-    def __init__(self, batch_transforms=None, collate_batch=True):
+    def __init__(self,
+                 batch_transforms=None,
+                 collate_batch=True,
+                 return_list=False):
         super(BatchCompose, self).__init__()
         self.batch_transforms = batch_transforms
         self.collate_batch = collate_batch
+        self.return_list = return_list
 
     def __call__(self, samples):
         if self.batch_transforms is not None:
@@ -51,7 +55,13 @@ class BatchCompose(Transform):
                 if k in sample:
                     sample.pop(k)
 
-        if self.collate_batch:
+        if self.return_list:
+            batch_data = [{
+                k: np.expand_dims(
+                    sample[k], axis=0)
+                for k in sample
+            } for sample in samples]
+        elif self.collate_batch:
             batch_data = default_collate_fn(samples)
         else:
             batch_data = {}

+ 25 - 11
paddlex/deploy.py

@@ -167,10 +167,16 @@ class Predictor(object):
                     'score_map': s
                 } for l, s in zip(label_map, score_map)]
         elif self._model.model_type == 'detector':
-            net_outputs = {
-                k: v
-                for k, v in zip(['bbox', 'bbox_num', 'mask'], net_outputs)
-            }
+            if 'RCNN' in self._model.__class__.__name__:
+                net_outputs = [{
+                    k: v
+                    for k, v in zip(['bbox', 'bbox_num', 'mask'], res)
+                } for res in net_outputs]
+            else:
+                net_outputs = {
+                    k: v
+                    for k, v in zip(['bbox', 'bbox_num', 'mask'], net_outputs)
+                }
             preds = self._model._postprocess(net_outputs)
             if len(preds) == 1:
                 preds = preds[0]
@@ -192,7 +198,6 @@ class Predictor(object):
             input_tensor = self.predictor.get_input_handle(name)
             input_tensor.copy_from_cpu(inputs[name])
 
-        self.timer.inference_time_s.start()
         self.predictor.run()
         output_names = self.predictor.get_output_names()
         net_outputs = list()
@@ -225,15 +230,24 @@ class Predictor(object):
         self.timer.preprocess_time_s.end()
 
         self.timer.inference_time_s.start()
-        net_outputs = self.raw_predict(preprocessed_input)
-        self.timer.inference_time_s.end()
+        if 'RCNN' in self._model.__class__.__name__:
+            if len(preprocessed_input) > 1:
+                logging.warning(
+                    "{} only supports inference with batch size equal to 1."
+                    .format(self._model.__class__.__name__))
+            net_outputs = [
+                self.raw_predict(sample) for sample in preprocessed_input
+            ]
+            self.timer.inference_time_s.end(repeats=len(preprocessed_input))
+            ori_shape = None
+        else:
+            net_outputs = self.raw_predict(preprocessed_input)
+            self.timer.inference_time_s.end()
+            ori_shape = preprocessed_input.get('ori_shape', None)
 
         self.timer.postprocess_time_s.start()
         results = self.postprocess(
-            net_outputs,
-            topk,
-            ori_shape=preprocessed_input.get('ori_shape', None),
-            transforms=transforms)
+            net_outputs, topk, ori_shape=ori_shape, transforms=transforms)
         self.timer.postprocess_time_s.end()
 
         self.timer.img_num = len(images)