12345678910111213141516171819202122232425262728293031323334353637383940 |
- 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("--------------------------------")
|