55 lines
1.5 KiB
Python
55 lines
1.5 KiB
Python
|
from PIL import Image
|
||
|
|
import cv2 as cv
|
||
|
|
import re
|
||
|
|
|
||
|
|
|
||
|
|
def detect_img(yolo, img_path):
|
||
|
|
detected_rats = []
|
||
|
|
|
||
|
|
try:
|
||
|
|
image = Image.open(img_path)
|
||
|
|
except:
|
||
|
|
print('Image open Error! Try again!')
|
||
|
|
return None
|
||
|
|
else:
|
||
|
|
r_image, pred = yolo.detect_image(image)
|
||
|
|
processed_image = cv.imread(img_path)
|
||
|
|
# r_image.show()
|
||
|
|
# r_image.save(img_path)
|
||
|
|
|
||
|
|
## FIXME : better list mapping
|
||
|
|
for prediction in pred:
|
||
|
|
is_rat_detected = re.search("rat", prediction[0])
|
||
|
|
if is_rat_detected:
|
||
|
|
x1 = prediction[1][0]
|
||
|
|
x2 = prediction[2][0]
|
||
|
|
y1 = prediction[1][1]
|
||
|
|
y2 = prediction[2][1]
|
||
|
|
w = abs(x1 - x2)
|
||
|
|
h = abs(y1 - y2)
|
||
|
|
# print(pred)
|
||
|
|
# print(f'x1: {x1}, x2: {x2}, y1: {y1}, y2: {y2}, w: {w}, h: {h}')
|
||
|
|
|
||
|
|
rat_img = processed_image[y1:y1 + h, x1:x1 + w]
|
||
|
|
rat_img = cv.resize(rat_img, (128, 128), interpolation=cv.INTER_AREA)
|
||
|
|
rat_pos = w, x1
|
||
|
|
detected_rats.append((rat_img, rat_pos))
|
||
|
|
|
||
|
|
return detected_rats
|
||
|
|
|
||
|
|
|
||
|
|
def get_turn_value(cords):
|
||
|
|
img_width = 1920
|
||
|
|
w, x = cords
|
||
|
|
center_of_object = (x + w) / 2
|
||
|
|
object_position = center_of_object / img_width
|
||
|
|
return round(object_position * 100, 2)
|
||
|
|
|
||
|
|
|
||
|
|
def detect_rat(model, img_path):
|
||
|
|
detected_rats = detect_img(model, img_path)
|
||
|
|
return detected_rats
|
||
|
|
# rat_position = detected_rats[0][1]
|
||
|
|
# turn_val = get_turn_value(rat_position)
|
||
|
|
# return turn_val
|