def __init__(self):
# Deploy darknet53 model on cooresponding device
yolov3 = darknet("cfg/yolov3-1.cfg", 1)
yolov3.load_weight("yolov3-1.weights")
# Deploy stacked hourglass model
stackedhourglass = StackedHourglass(16)
stackedhourglass.load_state_dict(torch.load("stacked_hourglass.pkl"))
cuda = torch.cuda.is_available()
if cuda:
yolov3.cuda()
stackedhourglass.cuda()
yolov3.eval()
self.detector = detector(yolov3)
self.estimator = Estimator(stackedhourglass)
def __init__(self):
# Deploy darknet53 model on cooresponding device
yolov3 = darknet("cfg/yolov3.cfg", 80)
yolov3.load_weight("yolov3.weights")
yolov3.eval()
# Deploy stacked hourglass model
stackedhourglass = demo.__dict__['hg'](num_stacks=2,
num_blocks=1,
num_classes=16)
stackedhourglass = torch.nn.DataParallel(stackedhourglass)
stackedhourglass.eval()
checkpoint = torch.load('demo/hg_s2_b1/model_best.pth.tar')
stackedhourglass.load_state_dict(checkpoint['state_dict'])
cuda = torch.cuda.is_available()
if cuda:
yolov3.cuda()
stackedhourglass.cuda()
self.detector = detector(yolov3)
self.estimator = stackedhourglass
def __init__(self,
minR,
maxR,
resR,
minTheta,
maxTheta,
resTheta,
minPhi,
maxPhi,
resPhi,
threshold,
mti=False):
'''
@description: Capture initialization. Walabot configuration is also completed here. Capture deploys a YOLO v3 detector to
detect human and once a person is detected and its central point is near the centre of optical image, sensors
will begin to collect data and save to predetermined directory.
@arges:
minR : (int) scan arena configuration parameter, minimum distance
maxR : (int) maximum distance of scan arena
resR : (float) resolution of depth
minTheta : (int) minimum theta
maxTheta : (int) maximum theta
resTheta : (int) vertical angular resolution
minPhi : (int) minimum phi
maxPhi : (int) maximum phi
resPhi : (int) horizontal angular resolution
threshold : (int) threshold for weak signals
mode : (string) scan mode
mti : (boolean) ignore static reflectors
'''
# YOLO v3 detector deployment
model = darknet("cfg/yolov3.cfg", 80)
model.load_weight("yolov3.weights")
model.cuda()
model.eval()
self.detector = detector(model)
'''
Args:
img : (ndarray) img matrix from cv2.imread(),
If you want to use plt.imread() or other
RGB format method, ensure to transform
from RGB to BGR, HWC to CHW
waitkey : (int) input for cv2.waitKey()
Returns:
Prediction bounding-boxes
'''
# Get input dimensions
# try:
prediction = self.detect(img)
for prediction_ in prediction:
coord1 = tuple(map(int, prediction_[:2]))
coord2 = tuple(map(int, prediction_[2:4]))
cv2.rectangle(img, coord1, coord2, (0, 255, 0), 2)
#finally:
cv2.imshow('prediction.jpg', img)
cv2.waitKey(waitkey)
if __name__ == "__main__":
model = darknet("D:/ShaoshuYang/HPE/cfg/yolov3-1.cfg", 80)
model.load_weight("src/yolov3-1-1.weights")
model.cuda()
test = detector(model)
img = cv2.imread("D:/ShaoshuYang/HPE/data/samples/sishui.jpg")
test.detect_test(img, 100000)