update preprocess

app.py

@@ -1,9 +1,9 @@
 import gradio as gr
 import time
-from onnx.counting import Counting
+from counting.counting import Counting
 
 
-counting = Counting("onnx/apgcc.onnx")
+counting = Counting("counting/apgcc.onnx")
 
 
 def filter_with_threshold(scores, points, threshold):
@@ -19,12 +19,14 @@ def filter_with_threshold(scores, points, threshold):
 def pred(img, threshold):
     # 计算处理时间
     start_at = time.time()
+    
+    processed_image, processed_image_original = counting.preprocess_image(img, True)
 
-    scores, points = counting.pred(img, is_rgb=True)
+    scores, points = counting.run_inference(processed_image)
 
     scores, points = filter_with_threshold(scores, points, threshold)
 
-    draw = counting.draw_pred(img, scores, points)
+    draw = counting.draw_pred(processed_image_original, scores, points)
 
     elapsed_time = time.time() - start_at
     use_time = f"use: {elapsed_time:.3f}s"

{onnx → counting}/counting.py

@@ -27,6 +27,8 @@ class Counting(BaseONNX):
             img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
         else:
             img = img
+            
+        img_copy = img.copy()
 
         # Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
         # 转换为 float32 类型
@@ -56,14 +58,14 @@ class Counting(BaseONNX):
 
         if h != new_h or w != new_w:
             img = cv2.resize(img, (new_w, new_h), interpolation=cv2.INTER_LINEAR)
-        
+            img_copy = cv2.resize(img_copy, (new_w, new_h), interpolation=cv2.INTER_LINEAR)
         # 调整维度顺序 (H,W,C) -> (C,H,W)
         img = np.transpose(img, (2, 0, 1))
         
         # 添加 batch 维度
         img = np.expand_dims(img, axis=0)
 
-        return img
+        return img, img_copy
 
     def run_inference(self, image: np.ndarray) -> any:
         """
@@ -97,7 +99,7 @@ class Counting(BaseONNX):
         else:
             img_bgr = image.copy()
 
-        processed_image = self.preprocess_image(img_bgr, is_rgb)
+        processed_image, _ = self.preprocess_image(img_bgr, is_rgb)
 
         scores, points = self.run_inference(processed_image)