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)