Python darknet.Darknet方法代码示例

# 需要导入模块: import darknet [as 别名]
# 或者: from darknet import Darknet [as 别名]
def load_model():
    CUDA = torch.cuda.is_available()
    classes = load_classes('data/coco.names')

    #Set up the neural network
    print("Loading network.....")
    model = Darknet(args.cfgfile)
    model.load_weights(args.weightsfile)
    print("Network successfully loaded")

    model.net_info["height"] = args.reso
    inp_dim = int(model.net_info["height"])
    assert inp_dim % 32 == 0
    assert inp_dim > 32

    #If there's a GPU availible, put the model on GPU
    if CUDA:
        model.cuda()

    #Set the model in evaluation mode
    model.eval()
    return model 

# 需要导入模块: import darknet [as 别名]
# 或者: from darknet import Darknet [as 别名]
def __init__(self, cfgfile, weightsfile):
        self.CONFIDENCE_THRESHOLD = 0.85
        self.NMS_THRESHOLD = 0.4
        self.NUM_CLASSES = 2 # hard code here for head detector
        self.CLASSES = [ 'Head' ]
        self.CUDA = torch.cuda.is_available()
        if self.CUDA:
            logging.info('Using CUDA.')
        else:
            logging.info('Using CPU.')
        logging.info("Loading network.....")
        self.model = Darknet(cfgfile)
        self.model.load_weights(weightsfile)
        self.model.net_info["height"] = 512 # hard code here because we didn't use Spp
        self.inp_dim = int(self.model.net_info["height"])
        if self.CUDA:
            self.model.cuda()
        self.model.eval()
        logging.info("Network successfully loaded.")

    # Detect the heads in the given image (opencv numpy array), and return the results 

# 需要导入模块: import darknet [as 别名]
# 或者: from darknet import Darknet [as 别名]
def main():

    args = parse_args()

    if args.cuda and not torch.cuda.is_available():
        print("ERROR: cuda is not available, try running on CPU")
        sys.exit(1)

    print('Loading network...')
    model = Darknet("cfg/yolov3.cfg")
    model.load_weights('yolov3.weights')
    if args.cuda:
        model.cuda()

    model.eval()
    print('Network loaded')

    if args.video:
        detect_video(model, args)

    else:
        detect_image(model, args) 

# 需要导入模块: import darknet [as 别名]
# 或者: from darknet import Darknet [as 别名]
def load_model():
    start = 0

    classes = load_classes(yolo_dir + '/data/coco.names')
    #Set up the neural network
    print("Loading YOLO network.....")
    model = Darknet(args.cfgfile)
    model.load_weights(args.weightsfile)
    print("Network successfully loaded")

    model.net_info["height"] = args.reso
    inp_dim = int(model.net_info["height"])
    assert inp_dim % 32 == 0
    assert inp_dim > 32

    if CUDA:
        model.cuda()

    model.eval()
    return model 

# 需要导入模块: import darknet [as 别名]
# 或者: from darknet import Darknet [as 别名]
def detect(cfgfile, weightfile, imgfile):
    m = Darknet(cfgfile)

    m.print_network()
    m.load_weights(weightfile)
    print('Loading weights from %s... Done!' % (weightfile))

    # if m.num_classes == 20:
    #     namesfile = 'data/voc.names'
    # elif m.num_classes == 80:
    #     namesfile = 'data/coco.names'
    # else:
    #     namesfile = 'data/names'
    
    use_cuda = torch.cuda.is_available()
    if use_cuda:
        m.cuda()

    img = Image.open(imgfile).convert('RGB')
    sized = letterbox_image(img, m.width, m.height)

    start = time.time()
    boxes = do_detect(m, sized, 0.5, 0.4, use_cuda)
    correct_yolo_boxes(boxes, img.width, img.height, m.width, m.height)

    finish = time.time()
    print('%s: Predicted in %f seconds.' % (imgfile, (finish-start)))

    class_names = load_class_names(namesfile)
    plot_boxes(img, boxes, 'predictions.jpg', class_names) 

# 需要导入模块: import darknet [as 别名]
# 或者: from darknet import Darknet [as 别名]
def detect_cv2(cfgfile, weightfile, imgfile):
    import cv2
    m = Darknet(cfgfile)

    m.print_network()
    m.load_weights(weightfile)
    print('Loading weights from %s... Done!' % (weightfile))

    if m.num_classes == 20:
        namesfile = 'data/voc.names'
    elif m.num_classes == 80:
        namesfile = 'data/coco.names'
    else:
        namesfile = 'data/names'
    
    use_cuda = True
    if use_cuda:
        m.cuda()

    img = cv2.imread(imgfile)
    sized = cv2.resize(img, (m.width, m.height))
    sized = cv2.cvtColor(sized, cv2.COLOR_BGR2RGB)
    
    for i in range(2):
        start = time.time()
        boxes = do_detect(m, sized, 0.5, 0.4, use_cuda)
        finish = time.time()
        if i == 1:
            print('%s: Predicted in %f seconds.' % (imgfile, (finish-start)))

    class_names = load_class_names(namesfile)
    plot_boxes_cv2(img, boxes, savename='predictions.jpg', class_names=class_names) 

# 需要导入模块: import darknet [as 别名]
# 或者: from darknet import Darknet [as 别名]
def detect_skimage(cfgfile, weightfile, imgfile):
    from skimage import io
    from skimage.transform import resize
    m = Darknet(cfgfile)

    m.print_network()
    m.load_weights(weightfile)
    print('Loading weights from %s... Done!' % (weightfile))

    if m.num_classes == 20:
        namesfile = 'data/voc.names'
    elif m.num_classes == 80:
        namesfile = 'data/coco.names'
    else:
        namesfile = 'data/names'
    
    use_cuda = True
    if use_cuda:
        m.cuda()

    img = io.imread(imgfile)
    sized = resize(img, (m.width, m.height)) * 255
    
    for i in range(2):
        start = time.time()
        boxes = do_detect(m, sized, 0.5, 0.4, use_cuda)
        finish = time.time()
        if i == 1:
            print('%s: Predicted in %f seconds.' % (imgfile, (finish-start)))

    class_names = load_class_names(namesfile)
    plot_boxes_cv2(img, boxes, savename='predictions.jpg', class_names=class_names) 

# 需要导入模块: import darknet [as 别名]
# 或者: from darknet import Darknet [as 别名]
def demo(cfgfile, weightfile):
    m = Darknet(cfgfile)
    m.print_network()
    m.load_weights(weightfile)
    print('Loading weights from %s... Done!' % (weightfile))

    if m.num_classes == 20:
        namesfile = 'data/voc.names'
    elif m.num_classes == 80:
        namesfile = 'data/coco.names'
    else:
        namesfile = 'data/names'
    print("{} is used for classification".format(namesfile))
    class_names = load_class_names(namesfile)
 
    use_cuda = True
    if use_cuda:
        m.cuda()

    cap = cv2.VideoCapture(1)
    if not cap.isOpened():
        print("Unable to open camera")
        exit(-1)

    while True:
        res, img = cap.read()
        if res:
            sized = cv2.resize(img, (m.width, m.height))
            bboxes = do_detect(m, sized, 0.5, 0.4, use_cuda)
            print('------')
            draw_img = plot_boxes_cv2(img, bboxes, None, class_names)
            cv2.imshow(cfgfile, draw_img)
            cv2.waitKey(1)
        else:
             print("Unable to read image")
             exit(-1) 

############################################ 

# 需要导入模块: import darknet [as 别名]
# 或者: from darknet import Darknet [as 别名]
def partial(cfgfile, weightfile, outfile, cutoff):
    m = Darknet(cfgfile)
    m.print_network()
    m.load_weights(weightfile)
    m.seen = 0
    m.save_weights(outfile, cutoff)
    print('save %s' % (outfile)) 

# 需要导入模块: import darknet [as 别名]
# 或者: from darknet import Darknet [as 别名]
def eval_widerface(cfgfile, weightfile, valdir, savedir):
    m = Darknet(cfgfile)
    m.load_weights(weightfile)
    use_cuda = 1
    if use_cuda:
        m.cuda()

    scale_size = 16
    class_names = load_class_names('data/names')
    for parent,dirnames,filenames in os.walk(valdir):
        if parent != valdir:
            targetdir = os.path.join(savedir, os.path.basename(parent))
            if not os.path.isdir(targetdir):
                os.mkdir(targetdir)
            for filename in filenames:
                imgfile = os.path.join(parent,filename)
                img = Image.open(imgfile).convert('RGB')
                sized_width = int(round(img.width*1.0/scale_size) * 16)
                sized_height = int(round(img.height*1.0/scale_size) * 16)
                sized = img.resize((sized_width, sized_height))
                print(filename, img.width, img.height, sized_width, sized_height)
                if sized_width * sized_height > 1024 * 2560:
                    print('omit %s' % filename)
                    continue
                boxes = do_detect(m, sized, 0.05, 0.4, use_cuda)
                if True:
                    savename = os.path.join(targetdir, filename)
                    print('save to %s' % savename)
                    plot_boxes(img, boxes, savename, class_names)
                if True:
                    savename = os.path.join(targetdir, os.path.splitext(filename)[0]+".txt")
                    print('save to %s' % savename)
                    save_boxes(img, boxes, savename) 

# 需要导入模块: import darknet [as 别名]
# 或者: from darknet import Darknet [as 别名]
def detect(cfgfile, weightfile, imgfolder):
    m = Darknet(cfgfile)

    #m.print_network()
    m.load_weights(weightfile)
    print('Loading weights from %s... Done!' % (weightfile))

    # if m.num_classes == 20:
    #     namesfile = 'data/voc.names'
    # elif m.num_classes == 80:
    #     namesfile = 'data/coco.names'
    # else:
    #     namesfile = 'data/names'
    
    use_cuda = True
    if use_cuda:
        m.cuda()

    imgfiles = [x for x in os.listdir(imgfolder) if x[-4:] == '.jpg']
    imgfiles.sort()
    for imgname in imgfiles:
        imgfile = os.path.join(imgfolder,imgname)
        
        img = Image.open(imgfile).convert('RGB')
        sized = img.resize((m.width, m.height))

        #for i in range(2):
        start = time.time()
        boxes = do_detect(m, sized, 0.5, 0.4, use_cuda)
        finish = time.time()
            #if i == 1:
        print('%s: Predicted in %f seconds.' % (imgfile, (finish-start)))

        class_names = load_class_names(namesfile)
        img = plot_boxes(img, boxes, 'result/{}'.format(os.path.basename(imgfile)), class_names)
        img = np.array(img)
        cv2.imshow('{}'.format(os.path.basename(imgfolder)), img)
        cv2.resizeWindow('{}'.format(os.path.basename(imgfolder)), 1000,800)
        cv2.waitKey(1000) 

# 需要导入模块: import darknet [as 别名]
# 或者: from darknet import Darknet [as 别名]
def __init__(self, data_options):
        super(SegPoseNet, self).__init__()

        pose_arch_cfg = data_options['pose_arch_cfg']
        self.width = int(data_options['width'])
        self.height = int(data_options['height'])
        self.channels = int(data_options['channels'])

        self.coreModel = Darknet(pose_arch_cfg, self.width, self.height, self.channels)
        self.segLayer = PoseSegLayer(data_options)
        self.regLayer = Pose2DLayer(data_options) 

# 需要导入模块: import darknet [as 别名]
# 或者: from darknet import Darknet [as 别名]
def load_model():
    scales = args.scales
    batch_size = int(args.bs)
    confidence = float(args.confidence)
    nms_thesh = float(args.nms_thresh)
    start = 0

    CUDA = torch.cuda.is_available()
    classes = load_classes(yolo_dir + '/data/coco.names')
    #Set up the neural network
    print("Loading YOLO network.....")
    model = Darknet(args.cfgfile)
    model.load_weights(args.weightsfile)
    print("Network successfully loaded")

    model.net_info["height"] = args.reso
    inp_dim = int(model.net_info["height"])
    assert inp_dim % 32 == 0
    assert inp_dim > 32

    #If there's a GPU availible, put the model on GPU
    if CUDA:
        model.cuda()

    #Set the model in evaluation mode
    model.eval()

    return model 

# 需要导入模块: import darknet [as 别名]
# 或者: from darknet import Darknet [as 别名]
def eval_list(cfgfile, namefile, weightfile, testfile):
    m = Darknet(cfgfile)
    m.load_weights(weightfile)
    use_cuda = 1
    if use_cuda:
        m.cuda()

    class_names = load_class_names(namefile)

    file_list = []
    with open(testfile, "r") as fin:
        for f in fin:
            file_list.append(f.strip())

    for imgfile in file_list:
        img = Image.open(imgfile).convert('RGB')
        sized = img.resize((m.width, m.height))
        filename = os.path.basename(imgfile)
        filename = os.path.splitext(filename)[0]
        #print(filename, img.width, img.height, sized_width, sized_height)

        if m.width * m.height > 1024 * 2560:
            print('omit %s' % filename)
            continue

        if False:
            boxes = do_detect(m, sized, conf_thresh, nms_thresh, use_cuda)
        else:
            m.eval()
            sized = image2torch(sized).cuda();
            #output = m(Variable(sized, volatile=True)).data
            output = m(sized)
            #boxes = get_region_boxes(output, conf_thresh, m.num_classes, m.anchors, m.num_anchors, 0, 1)[0]
            boxes = get_all_boxes(output, conf_thresh, m.num_classes)[0]
            boxes = np.array(nms(boxes, nms_thresh))

        if False:
            savename = get_det_image_name(imgfile)
            print('img: save to %s' % savename)
            plot_boxes(img, boxes, savename, class_names)

        if False:
            savename = get_det_result_name(imgfile)
            print('det: save to %s' % savename)
            save_boxes(imgfile, img, boxes, savename)