object-recognition-imageai/basic_object_detection_from_image.py at master · devopsbay/object-recognition-imageai · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
from datetime import datetime
import os

import cv2

from imageai.Detection import ObjectDetection
from imageai.Detection.Custom import CustomObjectDetection

probability = 30


def detect_objects_resnet(filename):

    detector = ObjectDetection()
    detector.setModelTypeAsRetinaNet()
    detector.setModelPath(os.path.join(os.getcwd(), "models/resnet50_coco_best_v2.0.1.h5"))
    detector.loadModel()

    detections = detector.detectObjectsFromImage(input_image=filename,
                                                 output_image_path=os.path.join('output/images/', datetime.now().strftime("%H_%M_%S")+"_resnet_detected.jpg"),
                                                 minimum_percentage_probability=probability)

    for eachObject in detections:
        print(eachObject["name"], " : ", eachObject["percentage_probability"], " : ", eachObject["box_points"])
        print("--------------------------------")
    # De-allocate any associated memory usage
    cv2.destroyAllWindows()


def detect_objects_yolo(filename):

    detector = ObjectDetection()
    detector.setModelTypeAsYOLOv3()
    detector.setModelPath(os.path.join(os.getcwd(), "models/yolo.h5"))
    detector.loadModel()
    detections = detector.detectObjectsFromImage(input_image=filename,
                                                 output_image_path=os.path.join('output/images/', datetime.now().strftime("%H_%M_%S")+"_yolo_detected.jpg"),
                                                 minimum_percentage_probability=probability,
                                                 extract_detected_objects=False)
    print("---------YOLO------------------")
    for eachObject in detections:
        print(eachObject["name"], " : ", eachObject["percentage_probability"], " : ", eachObject["box_points"])
        print("--------------------------------")
    # De-allocate any associated memory usage
    cv2.destroyAllWindows()


def detect_objects_yolo_custom(filename):

    detector = CustomObjectDetection()
    detector.setModelTypeAsYOLOv3()
    detector.setModelPath("examples/specjale/models/detection_model-ex-008--loss-0023.046.h5")
    detector.setJsonPath("examples/specjale/json/detection_config.json")
    detector.loadModel()

    detections = detector.detectObjectsFromImage(input_image=filename,
                                                 output_image_path=os.path.join('output/images/', datetime.now().strftime("%H_%M_%S")+"_yolo_detected.jpg"),
                                                 minimum_percentage_probability=probability,
                                                 extract_detected_objects=False)

    for eachObject in detections:
        print(eachObject["name"], " : ", eachObject["percentage_probability"], " : ", eachObject["box_points"])
        print("--------------------------------")
    # De-allocate any associated memory usage
    cv2.destroyAllWindows()


if __name__ == "__main__":
    filename = "examples/kuchnia4.jpg"
    #detect_objects_resnet(filename)
    #detect_objects_yolo(filename)
    detect_objects_yolo_custom(filename)