Reset git-lfs tracking

.gitattributes

@@ -33,3 +33,5 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+*.png filter=lfs diff=lfs merge=lfs -text
+*.jpg filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,4 @@
+coverage
+.DS_Store
+__pycache__
+*.pyc

app.py

@@ -0,0 +1,66 @@
+import gradio as gr
+import time
+from onnx.counting import Counting
+
+
+counting = Counting("onnx/apgcc.onnx")
+
+
+def filter_with_threshold(scores, points, threshold):
+    filtered_scores = []
+    filtered_points = []
+    for score, point in zip(scores, points):
+        if score > threshold:
+            filtered_scores.append(score)
+            filtered_points.append(point)
+    return filtered_scores, filtered_points
+
+
+def pred(img, threshold):
+    # 计算处理时间
+    start_at = time.time()
+
+    scores, points = counting.pred(img, is_rgb=True)
+
+    scores, points = filter_with_threshold(scores, points, threshold)
+
+    draw = counting.draw_pred(img, scores, points)
+
+    elapsed_time = time.time() - start_at
+    use_time = f"use: {elapsed_time:.3f}s"
+
+
+    total = len(points)
+
+    return draw, total, use_time
+
+
+model_description = """
+# APGCC People Counting
+
+APGCC (Adaptive Perspective Guidance for Crowd Counting) 
+
+### based on
+
+- [APGCC](https://github.com/AaronCIH/APGCC)
+"""
+
+demo = gr.Interface(
+    description=model_description,
+    fn=pred, 
+    inputs=["image", 
+            gr.Slider(0, 1, 0.5, label="Threshold")], 
+    outputs=[
+        "image",
+        gr.Number(label="Count"),
+        gr.Textbox(label="useTime"),
+    ],
+    examples=[
+        ["examples/crowd-001.jpg", 0.5],
+        ["examples/crowd-002.jpg", 0.5],
+        ["examples/image.png", 0.5],
+        ["examples/image2.png", 0.5],
+        ["examples/few-001.png", 0.5],
+    ])
+    
+demo.launch()

onnx/apgcc.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:85904b6137f2fd7f94fe6acc1d996f84d5a76dfeff904d10336584e5a4db68eb
+size 71659567

onnx/base_onnx.py

@@ -0,0 +1,87 @@
+import onnxruntime
+import numpy as np
+import cv2
+from abc import ABC, abstractmethod
+from typing import Any, Tuple, Union, List
+
+class BaseONNX(ABC):
+    def __init__(self, model_path: str):
+        """初始化ONNX模型基类
+
+        Args:
+            model_path (str): ONNX模型路径
+            input_size (tuple): 模型输入尺寸 (width, height)
+        """
+        self.session = onnxruntime.InferenceSession(model_path)
+        self.input_name = self.session.get_inputs()[0].name
+
+    def load_image(self, image: Union[cv2.UMat, str]) -> cv2.UMat:
+        """加载图像
+
+        Args:
+            image (Union[cv2.UMat, str]): 图像路径或cv2图像对象
+
+        Returns:
+            cv2.UMat: 加载的图像
+        """
+        if isinstance(image, str):
+            return cv2.imread(image)
+        return image.copy()
+
+    @abstractmethod
+    def preprocess_image(self, img_bgr: cv2.UMat, *args, **kwargs) -> np.ndarray:
+        """图像预处理抽象方法
+
+        Args:
+            img_bgr (cv2.UMat): BGR格式的输入图像
+            
+        Returns:
+            np.ndarray: 预处理后的图像
+        """
+        pass
+
+    @abstractmethod
+    def run_inference(self, image: np.ndarray) -> Any:
+        """运行推理的抽象方法
+
+        Args:
+            image (np.ndarray): 预处理后的输入图像
+
+        Returns:
+            Any: 模型输出结果
+        """
+        pass
+
+    @abstractmethod
+    def pred(self, image: Union[cv2.UMat, str], *args, **kwargs) -> Any:
+        """预测的抽象方法
+
+        Args:
+            image (Union[cv2.UMat, str]): 输入图像或图像路径
+
+        Returns:
+            Any: 预测结果
+        """
+        pass
+
+    @abstractmethod
+    def draw_pred(self, img: cv2.UMat, *args, **kwargs) -> cv2.UMat:
+        """绘制预测结果的抽象方法
+
+        Args:
+            img (cv2.UMat): 要绘制的图像
+
+        Returns:
+            cv2.UMat: 绘制结果后的图像
+        """
+        pass
+
+    
+    def check_images_list(self, images: List[Union[cv2.UMat, str, np.ndarray]]):
+        """
+        检查图像列表是否有效
+        """
+        for image in images:
+            if not isinstance(image, cv2.UMat) and not isinstance(image, str) and not isinstance(image, np.ndarray):
+                raise ValueError("The images must be a list of cv2.UMat or str or np.ndarray.")
+ 

onnx/counting.py

@@ -0,0 +1,119 @@
+import onnxruntime
+import numpy as np
+import cv2
+from typing import Tuple, List, Union
+from .base_onnx import BaseONNX
+
+class Counting(BaseONNX):
+    UPPER_BOUND = 2560 
+    MULTIPLE_OF = 32
+
+    def __init__(self,  model_path):
+        super().__init__(model_path)
+
+
+    def preprocess_image(self, img: cv2.UMat, is_rgb: bool = True):
+        """
+        预处理图像，包括颜色转换、缩放和标准化
+        
+        Args:
+            img: 输入图像，BGR或RGB格式
+            is_rgb: 是否已经是RGB格式，默认为True
+            
+        Returns:
+            预处理后的图像张量，形状为(1, 3, H, W)
+        """
+        if not is_rgb:
+            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+        else:
+            img = img
+
+        # Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+        # 转换为 float32 类型
+        img = img.astype(np.float32)
+        # 除以 255.0
+        img /= 255.0
+        # 减去均值
+        img -= np.array([0.485, 0.456, 0.406])
+        # 除以标准差
+        img /= np.array([0.229, 0.224, 0.225])
+
+
+        # 检查图像大小是否超过上限
+        origin_h, origin_w = img.shape[:2]
+        max_size = max(origin_h, origin_w)
+
+        if max_size > self.UPPER_BOUND:
+            scale = self.UPPER_BOUND / max_size
+            img = cv2.resize(img, None, fx=scale, fy=scale)
+        
+        
+        h, w = img.shape[:2]
+
+        # 确保图像尺寸是32的倍数
+        new_h = (h // self.MULTIPLE_OF) * self.MULTIPLE_OF
+        new_w = (w // self.MULTIPLE_OF) * self.MULTIPLE_OF
+
+        if h != new_h or w != new_w:
+            img = cv2.resize(img, (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
+
+    def run_inference(self, image: np.ndarray) -> any:
+        """
+        Run inference on the image.
+
+        Args:
+            image (np.ndarray): The image to run inference on.
+
+        Returns:
+            tuple: A tuple containing the detection results and labels.
+        """
+        
+        # 运行推理
+        result = self.session.run(None, {self.input_name: image})
+        
+        return result
+
+    def pred(self, image: List[Union[cv2.UMat, str]], is_rgb: bool = True) -> Tuple[List[float], List[List[float]],]:
+        """
+        Predict the detection results of the image.
+
+        Args:
+            image (cv2.UMat, str): The image to predict.
+
+        Returns:
+          
+        """
+        if isinstance(image, str):
+            img_bgr = cv2.imread(image)
+            is_rgb = False
+        else:
+            img_bgr = image.copy()
+
+        processed_image = self.preprocess_image(img_bgr, is_rgb)
+
+        scores, points = self.run_inference(processed_image)
+        
+        return scores, points
+    
+
+    def draw_pred(self, image: cv2.UMat, scores: List[float], points: List[List[float]]) -> cv2.UMat:
+
+        marked_img = np.array(image.copy())
+
+        for point, score in zip(points, scores):
+            # 确保点坐标在合理范围内
+            x, y = int(point[0]), int(point[1])
+            if 0 <= x < marked_img.shape[1] and 0 <= y < marked_img.shape[0]:
+                cv2.circle(marked_img, (x, y), 5, (255, 0, 0), -1)
+            
+        return marked_img
+
+

requirements.txt

@@ -0,0 +1,2 @@
+opencv-python
+onnxruntime