From 61fb2dbd20dd3085114e4bdf4efdecb89a4cdce2 Mon Sep 17 00:00:00 2001
From: Glenn Jocher <glenn.jocher@ultralytics.com>
Date: Wed, 9 Dec 2020 07:43:58 -0800
Subject: [PATCH] Simplify autoshape() post-process (#1603)

* Simplify autoshape() post-process

* cleanup
---
 hubconf.py       | 2 +-
 models/common.py | 7 +++----
 requirements.txt | 4 ++--
 utils/general.py | 2 +-
 4 files changed, 7 insertions(+), 8 deletions(-)

diff --git a/hubconf.py b/hubconf.py
index 9067d0bd..a0b66dc1 100644
--- a/hubconf.py
+++ b/hubconf.py
@@ -94,7 +94,7 @@ def yolov3_tiny(pretrained=False, channels=3, classes=80):
 
 if __name__ == '__main__':
     model = create(name='yolov3', pretrained=True, channels=3, classes=80)  # example
-    model = model.fuse().autoshape()  # for PIL/cv2/np inputs and NMS
+    model = model.autoshape()  # for PIL/cv2/np inputs and NMS
 
     # Verify inference
     from PIL import Image
diff --git a/models/common.py b/models/common.py
index 90edaab2..f26ffdd0 100644
--- a/models/common.py
+++ b/models/common.py
@@ -167,8 +167,7 @@ class autoShape(nn.Module):
 
         # Post-process
         for i in batch:
-            if y[i] is not None:
-                y[i][:, :4] = scale_coords(shape1, y[i][:, :4], shape0[i])
+            scale_coords(shape1, y[i][:, :4], shape0[i])
 
         return Detections(imgs, y, self.names)
 
@@ -177,13 +176,13 @@ class Detections:
     # detections class for YOLOv5 inference results
     def __init__(self, imgs, pred, names=None):
         super(Detections, self).__init__()
+        d = pred[0].device  # device
+        gn = [torch.tensor([*[im.shape[i] for i in [1, 0, 1, 0]], 1., 1.], device=d) for im in imgs]  # normalizations
         self.imgs = imgs  # list of images as numpy arrays
         self.pred = pred  # list of tensors pred[0] = (xyxy, conf, cls)
         self.names = names  # class names
         self.xyxy = pred  # xyxy pixels
         self.xywh = [xyxy2xywh(x) for x in pred]  # xywh pixels
-        d = pred[0].device  # device
-        gn = [torch.tensor([*[im.shape[i] for i in [1, 0, 1, 0]], 1., 1.], device=d) for im in imgs]  # normalizations
         self.xyxyn = [x / g for x, g in zip(self.xyxy, gn)]  # xyxy normalized
         self.xywhn = [x / g for x, g in zip(self.xywh, gn)]  # xywh normalized
         self.n = len(self.pred)
diff --git a/requirements.txt b/requirements.txt
index f9fc9fc2..4cb16138 100755
--- a/requirements.txt
+++ b/requirements.txt
@@ -26,5 +26,5 @@ pandas
 # scikit-learn==0.19.2  # for coreml quantization
 
 # extras --------------------------------------
-# thop  # FLOPS computation
-# pycocotools>=2.0  # COCO mAP
+thop  # FLOPS computation
+pycocotools>=2.0  # COCO mAP
diff --git a/utils/general.py b/utils/general.py
index aac0e44b..18285512 100755
--- a/utils/general.py
+++ b/utils/general.py
@@ -258,7 +258,7 @@ def wh_iou(wh1, wh2):
     return inter / (wh1.prod(2) + wh2.prod(2) - inter)  # iou = inter / (area1 + area2 - inter)
 
 
-def non_max_suppression(prediction, conf_thres=0.1, iou_thres=0.6, classes=None, agnostic=False, labels=()):
+def non_max_suppression(prediction, conf_thres=0.25, iou_thres=0.45, classes=None, agnostic=False, labels=()):
     """Performs Non-Maximum Suppression (NMS) on inference results
 
     Returns: