Spaces:
Runtime error
Runtime error
File size: 3,080 Bytes
6de956c |
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 |
import cv2 as cv
import numpy as np
from fastai.vision.all import *
import pathlib
def chop_img(img, step):
allImgs = []
width, height, i, j = 640, 512, 0, 0
for i in range(step):
for j in range(step):
# crop the image into step*rows and step*columns
imgCrop = img[int(0 + height / step * i): int(height / step + height / step * i),
int(0 + width / step * j): int(width / step + width / step * j)]
imgResize = cv.resize(imgCrop, (640, 512)) # Resize image
imgAug = augment(imgResize)
allImgs.append((imgResize, imgAug))
j += 1
i += 1
return allImgs
def augment(img):
kernel = np.ones((5, 5), np.uint8)
imgGray = cv.cvtColor(img, cv.COLOR_BGR2GRAY) # convert to grayscale
imgBlur = cv.GaussianBlur(imgGray, (5, 5), 0) # apply gaussian blur
imgThresh = cv.adaptiveThreshold( # apply adaptive threshold
imgBlur, 255, cv.ADAPTIVE_THRESH_GAUSSIAN_C, cv.THRESH_BINARY_INV, 7, 5)
imgDilation = cv.dilate(imgThresh, kernel, iterations=1) # apply dilation to amplify threshold result
return imgDilation
def crop_bat(img, box):
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])
# 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
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])
crop_x1 = top_left_x - 10
if crop_x1 <= 0:
crop_x1 = 1
crop_x2 = bot_right_x+11
if crop_x2 > 512:
crop_x2 = 512
crop_y1 = top_left_y-10
if crop_y1 <= 0:
crop_y1 = 1
crop_y2 = bot_right_y+11
if crop_y2 > 640:
crop_y2 = 640
bat_crop = img[crop_x1: crop_x2, crop_y1: crop_y2]
return bat_crop
def find_bats(allImgs):
batDepthMin, batDepthMax = 50, 400
cropped_bats = []
for img in allImgs:
blobs = cv.findContours(img[1], cv.RETR_LIST, cv.CHAIN_APPROX_SIMPLE)[-2]
# Process the blobs
for blob in blobs:
if batDepthMin < cv.contourArea(blob) < batDepthMax: # Only process blobs with a min / max size
rect = cv.minAreaRect(blob)
box = cv.boxPoints(rect)
box = np.int0(box)
cropped_bat = crop_bat(img[0], box)
cropped_bats.append(cropped_bat)
return cropped_bats
def bat_detector(img):
totalBats = 0
allImgs = chop_img(img, 3)
cropped_bats = find_bats(allImgs)
temp = pathlib.PosixPath
pathlib.PosixPath = pathlib.WindowsPath
learn = load_learner("model.pkl")
for bat in cropped_bats:
label, _, probs = learn.predict(bat)
p=f"{probs[0]:.4f}"
if label == '!bat' and p > '0.5':
pass
else:
totalBats += 1
return totalBats |