Skip to content

Commit f710522

Browse files
author
dujw
committed
fix yolov7 python
1 parent 64173e8 commit f710522

File tree

9 files changed

+49
-33
lines changed

9 files changed

+49
-33
lines changed

‎example-python/pytrt/__init__.py

Lines changed: 16 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -228,8 +228,14 @@ class ObjectBox(object):
228228

229229
class YoloType(Enum):
230230
V5 : int = 0
231-
V3 : int = 0
232231
X : int = 1
232+
V3 : int = 2
233+
V7 : int = 3
234+
235+
class ImageType(Enum):
236+
CVMat : int = 0
237+
GPUYUVNV12 : int = 1
238+
GPUBGR : int = 2
233239

234240
class NMSMethod(Enum):
235241
CPU : int = 0
@@ -283,6 +289,15 @@ def __init__(
283289
use_multi_preprocess_stream : bool = False
284290
): ...
285291
def commit(self, image : np.ndarray)->SharedFutureObjectBoxArray: ...
292+
def commit_gpu(
293+
self,
294+
imageptr : int, # GPU(device_id) data pointer
295+
width : int,
296+
height : int,
297+
device_id : int = 0,
298+
imtype : ImageType = ImageType.GPUBGR,
299+
stream : int = 0
300+
)->SharedFutureObjectBoxArray: ...
286301

287302
class CenterNet(object):
288303
valid : bool

‎example-python/test_yolov7.py

Lines changed: 28 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -1,19 +1,19 @@
11
import os
2-
import cv2
32
import numpy as np
43
import pytrt as tp
4+
import cv2
55

66
# change current workspace
77
os.chdir("../workspace/")
88

99
# 如果执行出错，请删掉 ~/.pytrt 的缓存模型
1010
# rm -rf ~/.pytrt，重新下载
11-
device_id = 0
11+
device_id = 2
1212
engine_file = "yolov7.fp32.trtmodel"
1313
if not os.path.exists(engine_file):
1414
tp.compile_onnx_to_file(5, tp.onnx_hub("yolov7"), engine_file, device_id=device_id)
1515

16-
yolo = tp.Yolo(engine_file, type=tp.YoloType.V5, device_id=device_id)
16+
yolo = tp.Yolo(engine_file, type=tp.YoloType.V7, device_id=device_id)
1717
image = cv2.imread("inference/car.jpg")
1818
bboxes = yolo.commit(image).get()
1919
print(f"{len(bboxes)} objects")
@@ -24,39 +24,38 @@
2424
cv2.rectangle(image, (left, top), (right, bottom), tp.random_color(box.class_label), 5)
2525

2626
os.makedirs("single_inference", exist_ok=True)
27-
saveto = "single_inference/yolov5.car.jpg"
27+
saveto = "single_inference/yolov7.car.jpg"
2828
print(f"Save to {saveto}")
2929

3030
cv2.imwrite(saveto, image)
3131

3232

33-
try:
34-
import torch
35-
36-
image = cv2.imread("inference/car.jpg")
37-
device = 0
33+
# try:
34+
# import torch # 如果出现invalid pointer错误，可以考虑把import torch放到import pytrt上面
35+
# image = cv2.imread("inference/car.jpg")
36+
# device = 0
3837

39-
gpu_image = torch.from_numpy(image).to(device)
40-
bboxes = yolo.commit_gpu(
41-
pimage = gpu_image.data_ptr(),
42-
width = image.shape[1],
43-
height = image.shape[0],
44-
device_id = device,
45-
imtype = tp.ImageType.GPUBGR,
46-
stream = torch.cuda.current_stream().cuda_stream
47-
).get()
48-
print(f"{len(bboxes)} objects")
49-
print(bboxes)
38+
# gpu_image = torch.from_numpy(image).to(device)
39+
# bboxes = yolo.commit_gpu(
40+
# pimage = gpu_image.data_ptr(),
41+
# width = image.shape[1],
42+
# height = image.shape[0],
43+
# device_id = device,
44+
# imtype = tp.ImageType.GPUBGR,
45+
# stream = torch.cuda.current_stream().cuda_stream
46+
# ).get()
47+
# print(f"{len(bboxes)} objects")
48+
# print(bboxes)
5049

51-
for box in bboxes:
52-
left, top, right, bottom = map(int, [box.left, box.top, box.right, box.bottom])
53-
cv2.rectangle(image, (left, top), (right, bottom), tp.random_color(box.class_label), 5)
50+
# for box in bboxes:
51+
# left, top, right, bottom = map(int, [box.left, box.top, box.right, box.bottom])
52+
# cv2.rectangle(image, (left, top), (right, bottom), tp.random_color(box.class_label), 5)
5453

55-
os.makedirs("single_inference", exist_ok=True)
56-
saveto = "single_inference/yolov7-gpuptr.car.jpg"
57-
print(f"Save to {saveto}")
54+
# os.makedirs("single_inference", exist_ok=True)
55+
# saveto = "single_inference/yolov7-gpuptr.car.jpg"
56+
# print(f"Save to {saveto}")
5857

59-
cv2.imwrite(saveto, image)
58+
# cv2.imwrite(saveto, image)
6059

61-
except Exception as e:
62-
print("GPUPtr test failed.", e)
60+
# except Exception as e:
61+
# print("GPUPtr test failed.", e)

‎example-simple_yolo/Makefile

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,5 @@
1+
cc := g++
2+
nvcc = ${lean_cuda}/bin/nvcc
13

24
cpp_srcs := $(shell find src -name "*.cpp")
35
cpp_objs := $(cpp_srcs:.cpp=.o)
@@ -64,7 +66,7 @@ link_librarys := $(foreach item,$(link_librarys),-l$(item))
6466
# 显卡对应的号码参考这里：https://developer.nvidia.com/zh-cn/cuda-gpus#compute
6567
# 如果是 jetson nano，提示找不到-m64指令，请删掉 -m64选项。不影响结果
6668
cpp_compile_flags := -std=c++11 -fPIC -m64 -g -fopenmp -w -O0
67-
cu_compile_flags := -std=c++11 -m64 -Xcompiler -fPIC -g -w -gencode=arch=compute_75,code=sm_75 -O0
69+
cu_compile_flags := -std=c++11 -m64 -Xcompiler -fPIC -g -w -O0
6870
link_flags := -pthread -fopenmp -Wl,-rpath='$$ORIGIN'
6971

7072
cpp_compile_flags += $(include_paths)
@@ -94,7 +96,7 @@ objs/%.o : src/%.cpp
9496
objs/%.cuo : src/%.cu
9597
@echo Compile CUDA $<
9698
@mkdir -p $(dir $@)
97-
@nvcc -c $< -o $@ $(cu_compile_flags)
99+
@${nvcc} -c $< -o $@ $(cu_compile_flags)
98100

99101
objs/%.mk : src/%.cpp
100102
@echo Compile depends CXX $<
@@ -104,7 +106,7 @@ objs/%.mk : src/%.cpp
104106
objs/%.cumk : src/%.cu
105107
@echo Compile depends CUDA $<
106108
@mkdir -p $(dir $@)
107-
@nvcc -M $< -MF $@ -MT $(@:.cumk=.cuo) $(cu_compile_flags)
109+
@${nvcc} -M $< -MF $@ -MT $(@:.cumk=.cuo) $(cu_compile_flags)
108110

109111
run : workspace/pro
110112
@cd workspace && ./pro

‎example-simple_yolo/workspace/simple_yolo_result/car.jpg

-478 KB
Binary file not shown.

‎example-simple_yolo/workspace/simple_yolo_result/gril.jpg

-77 KB
Binary file not shown.

‎example-simple_yolo/workspace/simple_yolo_result/group.jpg

-349 KB
Binary file not shown.

‎example-simple_yolo/workspace/simple_yolo_result/yq.jpg

-122 KB
Binary file not shown.

‎example-simple_yolo/workspace/simple_yolo_result/zand.jpg

-248 KB
Binary file not shown.

‎example-simple_yolo/workspace/simple_yolo_result/zgjr.jpg

-162 KB
Binary file not shown.

0 commit comments

Comments
 (0)