obj_detection/imageAi.py

import os

# import tensorflow as tf

# 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)

# gpus = tf.config.list_physical_devices('GPU')
# tf.config.experimental.set_virtual_device_configuration(
#     gpus[0],
#     [tf.config.experimental.VirtualDeviceConfiguration(memory_limit=1024)])

from imageai.Detection import ObjectDetection

execution_path = os.getcwd()
detector = ObjectDetection()
detector.setModelTypeAsYOLOv3()
# detector.setModelTypeAsRetinaNet()
detector.setModelPath(os.path.join(execution_path, 'yolo.h5'))
detector.loadModel()

detections = detector.detectObjectsFromImage(
    input_image=os.path.join(execution_path, 'circuit_board_img_result.png'),
    output_image_path=os.path.join(execution_path, 'snapshot_detected.jpg'),
    minimum_percentage_probability=30
)

print(detections, end='\n')
for eachObject in detections:
    print(eachObject['name'], ': ', eachObject['percentage_probability'])
    # print("--------------------------------")