classify.py 文件源码

def main(image_file):
    image = Image.open(image_file)
    if image is None:
        print 'Could not load image "%s"' % sys.argv[1]
        return

    image = np.array(image.convert('RGB'), dtype=np.uint8)
    image = image[:, :, ::-1].copy()

    winSize = (200, 200)
    stepSize = 32

    roi = extractRoi(image, winSize, stepSize)
    weight_map, mask_scale = next(roi)

    samples = [(rect, scale, cv2.cvtColor(window, cv2.COLOR_BGR2GRAY))
               for rect, scale, window in roi]

    X_test = [window for rect, scale, window in samples]
    coords = [(rect, scale) for rect, scale, window in samples]

    extractor = cv2.FeatureDetector_create('SURF')
    detector = cv2.DescriptorExtractor_create('SURF')

    affine = AffineInvariant(extractor, detector)

    saved = pickle.load(open('classifier.pkl', 'rb'))

    feature_transform = saved['pipe']
    model = saved['model']

    print 'Extracting Affine transform invariant features'
    affine_invariant_features = affine.transform(X_test)
    print 'Matching features with template'
    features = feature_transform.transform(affine_invariant_features)

    rects = classify(model, features, coords, weight_map, mask_scale)
    for (left, top, right, bottom) in non_max_suppression_fast(rects, 0.4):
        cv2.rectangle(image, (left, top), (right, bottom), (0, 0, 0), 10)
        cv2.rectangle(image, (left, top), (right, bottom), (32, 32, 255), 5)

    plt.imshow(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
    plt.show()