From 929c7f202251c6381d709ad6f2a3fc65860f3be9 Mon Sep 17 00:00:00 2001
From: Guanghua Yu <742925032@qq.com>
Date: Tue, 1 Sep 2020 19:27:55 +0800
Subject: [PATCH] support depoly infer for fcos (#1331)

---
 deploy/python/infer.py | 14 +++++++++++---
 1 file changed, 11 insertions(+), 3 deletions(-)

diff --git a/deploy/python/infer.py b/deploy/python/infer.py
index 9c9d934be..4f3c5ef21 100644
--- a/deploy/python/infer.py
+++ b/deploy/python/infer.py
@@ -74,7 +74,7 @@ class Resize(object):
         self.arch = arch
         self.use_cv2 = use_cv2
         self.interp = interp
-        self.scale_set = {'RCNN', 'RetinaNet'}
+        self.scale_set = {'RCNN', 'RetinaNet', 'FCOS'}
 
     def __call__(self, im, im_info):
         """
@@ -259,7 +259,7 @@ def create_inputs(im, im_info, model_arch='YOLO'):
         scale = scale_x
         im_info = np.array([resize_shape + [scale]]).astype('float32')
         inputs['im_info'] = im_info
-    elif 'RCNN' in model_arch:
+    elif ('RCNN' in model_arch) or ('FCOS' in model_arch):
         scale = scale_x
         im_info = np.array([resize_shape + [scale]]).astype('float32')
         im_shape = np.array([origin_shape + [1.]]).astype('float32')
@@ -276,7 +276,15 @@ class Config():
     Args:
         model_dir (str): root path of model.yml
     """
-    support_models = ['YOLO', 'SSD', 'RetinaNet', 'RCNN', 'Face', 'TTF']
+    support_models = [
+        'YOLO',
+        'SSD',
+        'RetinaNet',
+        'RCNN',
+        'Face',
+        'TTF',
+        'FCOS',
+    ]
 
     def __init__(self, model_dir):
         # parsing Yaml config for Preprocess
-- 
GitLab