mirror of
https://github.com/wassname/PSPNet-Keras-tensorflow.git
synced 2026-08-24 11:37:11 +08:00
49 lines
1.4 KiB
Python
49 lines
1.4 KiB
Python
import sys
|
|
|
|
SHARED_CAFFE_RESOLVER = None
|
|
|
|
class CaffeResolver(object):
|
|
def __init__(self):
|
|
self.import_caffe()
|
|
|
|
def import_caffe(self):
|
|
self.caffe = None
|
|
try:
|
|
# Try to import PyCaffe first
|
|
import caffe
|
|
self.caffe = caffe
|
|
except ImportError:
|
|
# Fall back to the protobuf implementation
|
|
from . import caffepb
|
|
self.caffepb = caffepb
|
|
show_fallback_warning()
|
|
if self.caffe:
|
|
# Use the protobuf code from the imported distribution.
|
|
# This way, Caffe variants with custom layers will work.
|
|
self.caffepb = self.caffe.proto.caffe_pb2
|
|
self.NetParameter = self.caffepb.NetParameter
|
|
|
|
def has_pycaffe(self):
|
|
return self.caffe is not None
|
|
|
|
def get_caffe_resolver():
|
|
global SHARED_CAFFE_RESOLVER
|
|
if SHARED_CAFFE_RESOLVER is None:
|
|
SHARED_CAFFE_RESOLVER = CaffeResolver()
|
|
return SHARED_CAFFE_RESOLVER
|
|
|
|
def has_pycaffe():
|
|
return get_caffe_resolver().has_pycaffe()
|
|
|
|
def show_fallback_warning():
|
|
msg = '''
|
|
------------------------------------------------------------
|
|
WARNING: PyCaffe not found!
|
|
Falling back to a pure protocol buffer implementation.
|
|
* Conversions will be drastically slower.
|
|
* This backend is UNTESTED!
|
|
------------------------------------------------------------
|
|
|
|
'''
|
|
sys.stderr.write(msg)
|