123456789101112131415161718192021222324252627282930313233343536 |
- '''Created Aug 15, 2021 deeejas'''
- import os
- # import tensorflow as tf
- # gpus = tf.config.list_physical_devices('GPU')
- # tf.config.experimental.set_virtual_device_configuration(
- # gpus[0],
- # [tf.config.experimental.VirtualDeviceConfiguration(memory_limit=1024)]
- # )
- # gpus = tf.config.experimental.list_physical_devices('GPU')
- # if gpus:
- # try:
- # # Currently, memory growth needs to be the same across GPUs
- # for gpu in gpus:
- # tf.config.experimental.set_memory_growth(gpu, True)
- # logical_gpus = tf.config.experimental.list_logical_devices('GPU')
- # print(len(gpus), "Physical GPUs,", len(logical_gpus), "Logical GPUs")
- # except RuntimeError as e:
- # # Memory growth must be set before GPUs have been initialized
- # print(e)
- from imageai.Classification import ImageClassification
- execution_path = os.getcwd()
- prediction = ImageClassification()
- prediction.setModelTypeAsMobileNetV2()
- prediction.setModelPath(os.path.join(execution_path, "mobilenet_v2.h5"))
- prediction.loadModel()
- predictions, probabilities = prediction.classifyImage(os.path.join(execution_path, "circuit_board_img_result.png"), result_count=5)
- for eachPrediction, eachProbability in zip(predictions, probabilities):
- print(eachPrediction, ": ", eachProbability)
|