-
Notifications
You must be signed in to change notification settings - Fork 9
/
Copy pathdemo1.py
89 lines (57 loc) · 2.05 KB
/
demo1.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
#-*- coding: UTF-8 -*-
#! /usr/bin/env python
"""
Reads Darknet53 config and weights and creates Caffe2 model
"""
import argparse
import configparser
import io
import os
import sys
import logging
import time
import yaml
from collections import defaultdict
import cv2
from caffe2.python import core
from caffe2.python import workspace
from caffe2.python import model_helper,brew
from caffe2.python import dyndep
from PIL import Image, ImageDraw
import numpy as np
logger = logging.getLogger(__name__)
yolo_ops_lib = myutils.get_yolo_ops_lib()
dyndep.InitOpsLibrary(yolo_ops_lib)
img_dir = '/opt/LabelMe2CocoDataset/decoded_img_48349.jpg'
model_dir = "/opt/caffe2_yolov3/yolov3.minidb"
def _main(args):
num_classes = 80
if num_classes == 20:
namesfile = 'data/voc.names'
elif num_classes == 80:
namesfile = 'data/coco.names'
else:
namesfile = 'data/names'
class_names = load_class_names(namesfile)
with myutils.NamedCudaScope(0):
workspace.ResetWorkspace()
predict_net = prepare_prediction_net(model_dir, "minidb")
print predict_net.Proto()
workspace.CreateNet(predict_net,overwrite=True)
print("The blobs in the workspace after loading the model: {}".format(workspace.Blobs()))
img = cv2.imread(img_dir)
b, g, r = cv2.split(img)
rgb_img = cv2.merge([r, g, b])
sized = cv2.resize(rgb_img, (416, 416), interpolation=cv2.INTER_CUBIC)
npar = np.array(sized)
pp = np.ascontiguousarray(np.transpose(npar,[2,0,1])).reshape(1,3,sized.shape[0],sized.shape[1]).astype(np.float32)/255.0
for i in range(1):
list_boxes = do_detect(predict_net,pp)
boxes = list_boxes[0] + list_boxes[1] + list_boxes[2]
boxes = nms(boxes, 0.4)
image = Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
plot_boxes(image, boxes, 'predictions.jpg', class_names)
if __name__ == '__main__':
core.GlobalInit(['caffe2', '--caffe2_log_level=0'])
workspace.ResetWorkspace()
_main(parser.parse_args())