添加qat量化支持

yolo11_standalone.py

@@ -562,6 +562,90 @@ class YOLO11E(YOLO11):
             
         return head(feats, cls_pe)
 
+
+# ==============================================================================
+# [Part 4] PostProcessorNumpy
+# ==============================================================================
+
+import numpy as np
+
+class YOLOPostProcessorNumpy:
+    def __init__(self, strides=[8, 16, 32], reg_max=16, use_segmentation=False):
+        self.strides = np.array(strides, dtype=np.float32)
+        self.reg_max = reg_max
+        self.use_segmentation = use_segmentation
+        self.anchors = None
+        self.strides_array = None
+        self.shape = None
+        self.dfl_weights = np.arange(reg_max, dtype=np.float32).reshape(1, 1, reg_max, 1)
+
+    def sigmoid(self, x):
+        return 1 / (1 + np.exp(-x))
+
+    def softmax(self, x, axis=-1):
+        x_max = np.max(x, axis=axis, keepdims=True)
+        e_x = np.exp(x - x_max)
+        return e_x / np.sum(e_x, axis=axis, keepdims=True)
+
+    def make_anchors(self, feats, strides, grid_cell_offset=0.5):
+        anchor_points, stride_list = [], []
+        for i, stride in enumerate(strides):
+            _, _, h, w = feats[i].shape
+            sx = np.arange(w, dtype=np.float32) + grid_cell_offset
+            sy = np.arange(h, dtype=np.float32) + grid_cell_offset
+            sy, sx = np.meshgrid(sy, sx, indexing='ij')
+            
+            anchor_points.append(np.stack((sx, sy), -1).reshape(-1, 2))
+            stride_list.append(np.full((h * w, 1), stride, dtype=np.float32))
+        
+        return np.concatenate(anchor_points), np.concatenate(stride_list)
+
+    def dist2bbox(self, distance, anchor_points, xywh=True, dim=-1):
+        lt, rb = np.split(distance, 2, axis=dim)
+        x1y1 = anchor_points - lt
+        x2y2 = anchor_points + rb
+        if xywh:
+            c_xy = (x1y1 + x2y2) / 2
+            wh = x2y2 - x1y1
+            return np.concatenate((c_xy, wh), axis=dim)
+        return np.concatenate((x1y1, x2y2), axis=dim)
+
+    def dfl_decode(self, x):
+        B, C, A = x.shape
+        x = x.reshape(B, 4, self.reg_max, A)
+        x = self.softmax(x, axis=2)
+        return np.sum(x * self.dfl_weights, axis=2)
+
+    def __call__(self, outputs):
+        if self.use_segmentation:
+            x, mc, p = outputs
+        else:
+            x = outputs 
+
+        current_no = x[0].shape[1] 
+        current_nc = current_no - self.reg_max * 4
+        shape = x[0].shape
+        
+        x_cat = np.concatenate([xi.reshape(shape[0], current_no, -1) for xi in x], axis=2)
+        
+        if self.anchors is None or self.shape != shape:
+            self.anchors, self.strides_array = self.make_anchors(x, self.strides, 0.5)
+            self.shape = shape
+
+        box, cls = np.split(x_cat, [self.reg_max * 4], axis=1)  
+        dist = self.dfl_decode(box)
+        dist = dist.transpose(0, 2, 1) 
+        dbox = self.dist2bbox(dist, self.anchors, xywh=True, dim=2) * self.strides_array
+        cls = cls.transpose(0, 2, 1)
+        sigmoid_cls = self.sigmoid(cls)
+        final_box = np.concatenate((dbox, sigmoid_cls), axis=2)
+        
+        if self.use_segmentation:
+            return final_box, mc, p
+            
+        return final_box
+    
+
 if __name__ == "__main__":
     print("Testing Standard YOLO11...")
     model_std = YOLO11(nc=80, scale='n')