forked from cchighman/PriusWatchML
-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathPriusObjectDetection.py
104 lines (80 loc) · 4.08 KB
/
PriusObjectDetection.py
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
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
import multiprocessing
import os
import queue
import threading
import time
from datetime import timedelta
from multiprocessing import Pool
import numpy as np
from PIL import Image
from PriusImage import PriusImage
from PriusPalette import PriusPalette
from imageai.Detection import ObjectDetection
from imageai.Prediction.Custom import CustomImagePrediction
class PriusPredictor(object):
def __init__(self, image_path, model_path, output_path):
self.avgColor = []
self.pcaColors = []
self.detector = ObjectDetection()
self.detector.setModelTypeAsYOLOv3()
self.detector.setModelPath(model_path + "yolo.h5")
self.detector.loadModel(detection_speed="flash")
self.prediction = CustomImagePrediction()
self.prediction.setModelTypeAsResNet()
#self.prediction.setModelPath(model_path + "model_ex-012_acc-0.988819.h5")
self.prediction.setModelPath(model_path + "model_ex-043_acc-0.996787.h5")
self.prediction.setJsonPath(model_path + "model_class.json")
self.prediction.loadModel(num_objects=2)
now = time.localtime()
self.frame_folder = str(now.tm_year) + str(now.tm_mon) + str(now.tm_mday)
self.image_path = image_path
self.output_path = output_path + "detection/" + self.frame_folder + "/"
if os.path.exists(image_path) is False:
os.mkdir(image_path)
if os.path.exists(output_path) is False:
os.mkdir(output_path)
if os.path.exists(os.path.join(output_path, 'detection')) is False:
os.mkdir(os.path.join(output_path, 'detection'))
if os.path.exists(os.path.join(output_path, 'processed')) is False:
os.mkdir(os.path.join(output_path, 'processed'))
self.create_output_folder()
def create_output_folder(self):
if os.path.exists(self.output_path) is False:
os.mkdir(self.output_path)
def predict_vehicle_method(self, prediction_meta):
detected_img = os.path.join(prediction_meta['image_path'], prediction_meta['image_name'])
if os.path.exists(detected_img) is not True:
detected_img = prediction_meta['image_path']
return self.prediction.predictImage(detected_img, result_count=2)
def predict_vehicle(self, prediction_meta):
detected_img = prediction_meta['image_path']
return self.prediction.predictImage(detected_img, result_count=2)
def detect_pca(self, image):
priusImage = PriusImage.from_path(image)
return priusImage.has_pca_match()
def detect_vehicle(self, meta_data):
try:
image = os.path.join(meta_data["image_path"], meta_data['image_name'])
output_image = self.output_path + meta_data['image_name']
#print("Detecting vehicle for " + meta_data['image_name'] + " -> " + output_image)
if os.path.exists(image) is not True:
print("File doesnt exist. File: " + image)
custom_objects = self.detector.CustomObjects(car=True)
detections, objects_path = self.detector.detectCustomObjectsFromImage(custom_objects=custom_objects,
input_image=image,
extract_detected_objects=True,
output_image_path=output_image,
minimum_percentage_probability=50)
return zip(detections, objects_path)
except Exception as e:
print("While detecting vehicle: " + str(e))
def detect_vehicle_from_array(self, decoded):
#print("Detecting vehicle for " + meta_data['image_name'] + " -> " + output_image)
custom_objects = self.detector.CustomObjects(car=True)
result = self.detector.detectCustomObjectsFromImage(custom_objects=custom_objects,
input_type="array",
input_image=np.array(decoded),
#output_type="array",
minimum_percentage_probability=50)
print("Detected: " + str(result))
return result