feat: TensorRT 固定 batch=4 重构
- tensorrt_engine.py 工业级 Buffer Pool - preprocessor.py 添加 pad_to_batch4() - postprocessor.py 支持批量输出 - settings.py 固定 batch_size=4
This commit is contained in:
@@ -597,20 +597,21 @@ class PostProcessor:
|
||||
|
||||
output = outputs[0]
|
||||
|
||||
if len(output.shape) == 3:
|
||||
if output.ndim == 3:
|
||||
output = output[0]
|
||||
|
||||
num_detections = output.shape[0]
|
||||
if output.ndim == 2:
|
||||
output = output.reshape(-1)
|
||||
|
||||
num_detections = output.shape[0] // 85
|
||||
|
||||
boxes = []
|
||||
scores = []
|
||||
class_ids = []
|
||||
|
||||
for i in range(num_detections):
|
||||
detection = output[i]
|
||||
|
||||
if len(detection) < 6:
|
||||
continue
|
||||
start_idx = i * 85
|
||||
detection = output[start_idx:start_idx + 85]
|
||||
|
||||
x_center = detection[0]
|
||||
y_center = detection[1]
|
||||
@@ -637,16 +638,16 @@ class PostProcessor:
|
||||
y2 = y_center + height / 2
|
||||
|
||||
boxes.append([x1, y1, x2, y2])
|
||||
scores.append(total_conf)
|
||||
class_ids.append(class_id)
|
||||
scores.append(float(total_conf))
|
||||
class_ids.append(int(class_id))
|
||||
|
||||
if not boxes:
|
||||
return np.array([]), np.array([]), np.array([])
|
||||
|
||||
return (
|
||||
np.array(boxes),
|
||||
np.array(scores),
|
||||
np.array(class_ids)
|
||||
np.array(boxes, dtype=np.float32),
|
||||
np.array(scores, dtype=np.float32),
|
||||
np.array(class_ids, dtype=np.int32)
|
||||
)
|
||||
|
||||
def filter_by_roi(
|
||||
|
||||
Reference in New Issue
Block a user