jonathan-roos commited on
Commit
6de956c
1 Parent(s): 77e9330

added detector.py

Browse files
Files changed (1) hide show
  1. detector.py +92 -0
detector.py ADDED
@@ -0,0 +1,92 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import cv2 as cv
2
+ import numpy as np
3
+ from fastai.vision.all import *
4
+ import pathlib
5
+
6
+ def chop_img(img, step):
7
+ allImgs = []
8
+ width, height, i, j = 640, 512, 0, 0
9
+ for i in range(step):
10
+ for j in range(step):
11
+ # crop the image into step*rows and step*columns
12
+ imgCrop = img[int(0 + height / step * i): int(height / step + height / step * i),
13
+ int(0 + width / step * j): int(width / step + width / step * j)]
14
+ imgResize = cv.resize(imgCrop, (640, 512)) # Resize image
15
+ imgAug = augment(imgResize)
16
+ allImgs.append((imgResize, imgAug))
17
+ j += 1
18
+ i += 1
19
+ return allImgs
20
+
21
+ def augment(img):
22
+ kernel = np.ones((5, 5), np.uint8)
23
+ imgGray = cv.cvtColor(img, cv.COLOR_BGR2GRAY) # convert to grayscale
24
+ imgBlur = cv.GaussianBlur(imgGray, (5, 5), 0) # apply gaussian blur
25
+ imgThresh = cv.adaptiveThreshold( # apply adaptive threshold
26
+ imgBlur, 255, cv.ADAPTIVE_THRESH_GAUSSIAN_C, cv.THRESH_BINARY_INV, 7, 5)
27
+ imgDilation = cv.dilate(imgThresh, kernel, iterations=1) # apply dilation to amplify threshold result
28
+ return imgDilation
29
+
30
+ def crop_bat(img, box):
31
+ x1, y1, x2, y2, x3, y3, x4, y4 = int(box[0][1]), int(box[0][0]), int(box[1][1]), int(box[1][0]), int(box[2][1]), int(box[2][0]), int(box[3][1]), int(box[3][0])
32
+
33
+ # Find distance from cx to top_left_x and cy to top_left_y to determine how many pixels the border around the cropped image should be
34
+ top_left_x, top_left_y, bot_right_x, bot_right_y = min([x1,x2,x3,x4]), min([y1,y2,y3,y4]), max([x1,x2,x3,x4]), max([y1,y2,y3,y4])
35
+
36
+ crop_x1 = top_left_x - 10
37
+ if crop_x1 <= 0:
38
+ crop_x1 = 1
39
+
40
+ crop_x2 = bot_right_x+11
41
+ if crop_x2 > 512:
42
+ crop_x2 = 512
43
+
44
+ crop_y1 = top_left_y-10
45
+ if crop_y1 <= 0:
46
+ crop_y1 = 1
47
+
48
+ crop_y2 = bot_right_y+11
49
+ if crop_y2 > 640:
50
+ crop_y2 = 640
51
+
52
+ bat_crop = img[crop_x1: crop_x2, crop_y1: crop_y2]
53
+
54
+ return bat_crop
55
+
56
+ def find_bats(allImgs):
57
+ batDepthMin, batDepthMax = 50, 400
58
+ cropped_bats = []
59
+
60
+ for img in allImgs:
61
+ blobs = cv.findContours(img[1], cv.RETR_LIST, cv.CHAIN_APPROX_SIMPLE)[-2]
62
+ # Process the blobs
63
+ for blob in blobs:
64
+ if batDepthMin < cv.contourArea(blob) < batDepthMax: # Only process blobs with a min / max size
65
+ rect = cv.minAreaRect(blob)
66
+ box = cv.boxPoints(rect)
67
+ box = np.int0(box)
68
+ cropped_bat = crop_bat(img[0], box)
69
+ cropped_bats.append(cropped_bat)
70
+
71
+ return cropped_bats
72
+
73
+ def bat_detector(img):
74
+ img = cv.resize(img, (640,512))
75
+ totalBats = 0
76
+
77
+ allImgs = chop_img(img, 3)
78
+
79
+ cropped_bats = find_bats(allImgs)
80
+
81
+ temp = pathlib.PosixPath
82
+ pathlib.PosixPath = pathlib.WindowsPath
83
+ learn = load_learner("model.pkl")
84
+
85
+ for bat in cropped_bats:
86
+ label, _, probs = learn.predict(bat)
87
+ p=f"{probs[0]:.4f}"
88
+ if label == '!bat' and p > '0.5':
89
+ pass
90
+ else:
91
+ totalBats += 1
92
+ return totalBats