/
dogDetector.py
106 lines (88 loc) · 3.12 KB
/
dogDetector.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
105
106
from darkflow.net.build import TFNet
from common.singleton import BaseClassSingleton
import dlib
import os
import numpy as np
root_path = os.path.dirname(os.path.abspath(__file__))
class DogDetector():
def __init__(self, cfg='yolo', weights='yolov2', threshold=0.3):
model_path = os.path.join(root_path, 'cfg/' + cfg + '.cfg')
load_path = os.path.join(root_path, 'weights/' + weights + '.weights')
self.options = {
'model': model_path,
'load': load_path,
'threshold': threshold,
'gpu': 1.0,
'savepb': True
}
print(self.options['model'], self.options['load'])
self.tfnet = TFNet(self.options)
def detectsOneDog(self, img):
result = self.tfnet.return_predict(img)
dog_list = []
for res in result:
if res['label'] == 'dog':
dog_list.append(res)
if len(dog_list) == 0:
return False
else:
print(dog_list)
temp_area = 0
index = 0
for i in range(len(dog_list)):
rect = self.getDogRect(dog_list[i], img)
area = rect[2] * rect[3]
if area > temp_area:
temp_area = area
index = i
print(dog_list[index])
return dog_list[index]
def detectDogHead(self, img):
detector = dlib.simple_object_detector(os.path.join(root_path, "data/dog_detector.svm"))
dets = detector(img)
if(len(dets) == 0):
print("Can't detect Head")
return False
else :
dets_pop = dets.pop()
return dets_pop
def isLeft(self, dog_rect, head_rect):
if(head_rect[0] + head_rect[2] > dog_rect[0]+dog_rect[2]):
return False;
else:
return True;
def getDogRect(self, result, originalImg):
tl = result['topleft']
br = result['bottomright']
width = br['x'] - tl['x']
height = br['y'] - tl['y']
error_pixel = round(max(width, height) / 4)
width += error_pixel
height += error_pixel
x = max(tl['x'] - round(error_pixel / 2), 1)
y = max(tl['y'] - round(error_pixel / 2), 1)
imgHeight, imgWidth = originalImg.shape[:2]
if (x + width) >= imgWidth:
width = imgWidth - x - 1
if (y + height) >= imgHeight:
height = imgHeight - y - 1
return (x, y, width, height)
def getDogPartRect(self, result, originalImg):
t = result.top()
l = result.left()
b = result.bottom()
r = result.right()
print(t, l, b, r)
width = r - l
height = b - t
error_pixel = round(max(width, height) / 4)
width += error_pixel
height += error_pixel
x = max(l - round(error_pixel / 2), 0)
y = max(t - round(error_pixel / 2), 0)
imgHeight, imgWidth = originalImg.shape[:2]
if (x + width) >= imgWidth:
width = imgWidth - x - 1
if (y + height) >= imgHeight:
height = imgHeight - y - 1
return (x, y, width, height)