luanhailiang
/
billiard
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

device_count

main
luanhailiang 2023-09-22 20:50:12 +08:00
parent 07c5e58149
commit 05bdd0abfe
1 changed files with 4 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
app.py
View File

@ -1,6 +1,7 @@
import cv2 import cv2
import json import json
import time import time
import torch
import requests import requests
import threading import threading
import numpy as np import numpy as np
@ -61,6 +62,8 @@ def worker(id,source,region=None,stream=False):
region = 10 region = 10
else: else:
region = int(region) region = int(region)
device_count = torch.cuda.device_count()
print("GPU device count:", device_count)
print("start loading model...",id,source,region) print("start loading model...",id,source,region)
model = YOLO('./best.pt') model = YOLO('./best.pt')
print("start loaded model!!!") print("start loaded model!!!")
@ -92,7 +95,7 @@ def worker(id,source,region=None,stream=False):
print("Error opening video stream....") print("Error opening video stream....")
break break
continue continue
result = model.track(frame,show=False,stream=False,persist=True,device=int(id)%4) result = model.track(frame,show=False,stream=False,persist=True,device=int(id)%device_count)
use += " track:"+str((time.time()*1000) - bgn) use += " track:"+str((time.time()*1000) - bgn)
del(ret) del(ret)
del(frame) del(frame)