Add TensorFlow and TFLite export

detect.py

@@ -5,6 +5,9 @@
 import cv2
 import torch
 import torch.backends.cudnn as cudnn
+import yaml
+from numpy import random
+import numpy as np
 
 from models.experimental import attempt_load
 from utils.datasets import LoadStreams, LoadImages
@@ -31,12 +34,57 @@ def detect(opt):
     half = device.type != 'cpu'  # half precision only supported on CUDA
 
     # Load model
-    model = attempt_load(weights, map_location=device)  # load FP32 model
-    stride = int(model.stride.max())  # model stride
-    imgsz = check_img_size(imgsz, s=stride)  # check img_size
-    names = model.module.names if hasattr(model, 'module') else model.names  # get class names
-    if half:
-        model.half()  # to FP16
+    weights = weights[0] if isinstance(weights, list) else weights
+    suffix = Path(weights).suffix
+    if suffix == '.pt':
+        backend = 'pytorch'
+        model = attempt_load(weights, map_location=device)  # load FP32 model
+        stride = int(model.stride.max())  # model stride
+        imgsz = check_img_size(imgsz, s=stride)  # check img_size
+        names = model.module.names if hasattr(model, 'module') else model.names  # class names
+        if half:
+            model.half()  # to FP16
+    else:
+        import tensorflow as tf
+        from tensorflow import keras
+
+        stride = None
+        with open('data/coco.yaml') as f:
+            names = yaml.load(f, Loader=yaml.FullLoader)['names']  # class names (assume COCO)
+
+        if suffix == '.pb':
+            backend = 'graph_def'
+
+            # https://www.tensorflow.org/guide/migrate#a_graphpb_or_graphpbtxt
+            # https://github.com/leimao/Frozen_Graph_TensorFlow
+            def wrap_frozen_graph(graph_def, inputs, outputs):
+                def _imports_graph_def():
+                    tf.compat.v1.import_graph_def(graph_def, name="")
+
+                wrapped_import = tf.compat.v1.wrap_function(_imports_graph_def, [])
+                import_graph = wrapped_import.graph
+                return wrapped_import.prune(
+                    tf.nest.map_structure(import_graph.as_graph_element, inputs),
+                    tf.nest.map_structure(import_graph.as_graph_element, outputs))
+
+            graph = tf.Graph()
+            graph_def = graph.as_graph_def()
+            graph_def.ParseFromString(open(weights, 'rb').read())
+            frozen_func = wrap_frozen_graph(graph_def=graph_def, inputs="x:0", outputs="Identity:0")
+
+        elif suffix == '.tflite':
+            backend = 'tflite'
+            # Load TFLite model and allocate tensors
+            interpreter = tf.lite.Interpreter(model_path=weights)
+            interpreter.allocate_tensors()
+
+            # Get input and output tensors
+            input_details = interpreter.get_input_details()
+            output_details = interpreter.get_output_details()
+
+        else:
+            backend = 'saved_model'
+            model = keras.models.load_model(weights)
 
     # Second-stage classifier
     classify = False
@@ -49,24 +97,49 @@ def detect(opt):
     if webcam:
         view_img = check_imshow()
         cudnn.benchmark = True  # set True to speed up constant image size inference
-        dataset = LoadStreams(source, img_size=imgsz, stride=stride)
+        dataset = LoadStreams(source, img_size=imgsz, stride=stride, auto=backend == 'pytorch')
     else:
-        dataset = LoadImages(source, img_size=imgsz, stride=stride)
+        dataset = LoadImages(source, img_size=imgsz, stride=stride, auto=backend == 'pytorch')
 
     # Run inference
-    if device.type != 'cpu':
-        model(torch.zeros(1, 3, imgsz, imgsz).to(device).type_as(next(model.parameters())))  # run once
+    if device.type != 'cpu' and backend == 'pytorch':
+        model(torch.zeros(1, 3, *imgsz).to(device).type_as(next(model.parameters())))  # run once
     t0 = time.time()
     for path, img, im0s, vid_cap in dataset:
         img = torch.from_numpy(img).to(device)
-        img = img.half() if half else img.float()  # uint8 to fp16/32
+        img = img.half() if half and backend == 'pytorch' else img.float()  # uint8 to fp16/32
         img /= 255.0  # 0 - 255 to 0.0 - 1.0
         if img.ndimension() == 3:
             img = img.unsqueeze(0)
 
         # Inference
         t1 = time_synchronized()
-        pred = model(img, augment=opt.augment)[0]
+        if backend == 'pytorch':
+            pred = model(img, augment=opt.augment)[0]
+        else:
+            if backend == 'saved_model':
+                pred = model(img.permute(0, 2, 3, 1).cpu().numpy(), training=False).numpy()
+            elif backend == 'graph_def':
+                pred = frozen_func(x=tf.constant(img.permute(0, 2, 3, 1).cpu().numpy())).numpy()
+            elif backend == 'tflite':
+                input_data = img.permute(0, 2, 3, 1).cpu().numpy()
+                if opt.tfl_int8:
+                    scale, zero_point = input_details[0]['quantization']
+                    input_data = input_data / scale + zero_point
+                    input_data = input_data.astype(np.uint8)
+                interpreter.set_tensor(input_details[0]['index'], input_data)
+                interpreter.invoke()
+                pred = interpreter.get_tensor(output_details[0]['index'])
+                if opt.tfl_int8:
+                    scale, zero_point = output_details[0]['quantization']
+                    pred = pred.astype(np.float32)
+                    pred = (pred - zero_point) * scale
+            # Denormalize xywh
+            pred[..., 0] *= imgsz[1]  # x
+            pred[..., 1] *= imgsz[0]  # y
+            pred[..., 2] *= imgsz[1]  # w
+            pred[..., 3] *= imgsz[0]  # h
+            pred = torch.tensor(pred)
 
         # Apply NMS
         pred = non_max_suppression(pred, opt.conf_thres, opt.iou_thres, opt.classes, opt.agnostic_nms,
@@ -152,7 +225,7 @@ def detect(opt):
     parser = argparse.ArgumentParser()
     parser.add_argument('--weights', nargs='+', type=str, default='yolov5s.pt', help='model.pt path(s)')
     parser.add_argument('--source', type=str, default='data/images', help='source')  # file/folder, 0 for webcam
-    parser.add_argument('--img-size', type=int, default=640, help='inference size (pixels)')
+    parser.add_argument('--img-size', nargs='+', type=int, default=[320, 320], help='image size')  # height, width
     parser.add_argument('--conf-thres', type=float, default=0.25, help='object confidence threshold')
     parser.add_argument('--iou-thres', type=float, default=0.45, help='IOU threshold for NMS')
     parser.add_argument('--max-det', type=int, default=1000, help='maximum number of detections per image')
@@ -169,10 +242,12 @@ def detect(opt):
     parser.add_argument('--project', default='runs/detect', help='save results to project/name')
     parser.add_argument('--name', default='exp', help='save results to project/name')
     parser.add_argument('--exist-ok', action='store_true', help='existing project/name ok, do not increment')
+    parser.add_argument('--tfl-int8', action='store_true', help='use int8 quantized TFLite model')
     parser.add_argument('--line-thickness', default=3, type=int, help='bounding box thickness (pixels)')
     parser.add_argument('--hide-labels', default=False, action='store_true', help='hide labels')
     parser.add_argument('--hide-conf', default=False, action='store_true', help='hide confidences')
     opt = parser.parse_args()
+    opt.img_size *= 2 if len(opt.img_size) == 1 else 1  # expand
     print(opt)
     check_requirements(exclude=('tensorboard', 'pycocotools', 'thop'))
 

models/experimental.py

@@ -110,14 +110,18 @@ def forward(self, x, augment=False):
         return y, None  # inference, train output
 
 
-def attempt_load(weights, map_location=None, inplace=True):
+def attempt_load(weights, map_location=None, inplace=True, fuse=True):
     from models.yolo import Detect, Model
 
     # Loads an ensemble of models weights=[a,b,c] or a single model weights=[a] or weights=a
     model = Ensemble()
     for w in weights if isinstance(weights, list) else [weights]:
         ckpt = torch.load(attempt_download(w), map_location=map_location)  # load
-        model.append(ckpt['ema' if ckpt.get('ema') else 'model'].float().fuse().eval())  # FP32 model
+        if fuse:
+            model.append(ckpt['ema' if ckpt.get('ema') else 'model'].float().fuse().eval())  # FP32 model
+        else:
+            model.append(ckpt['ema' if ckpt.get('ema') else 'model'].float().eval())  # without layer fuse
+
 
     # Compatibility updates
     for m in model.modules():