Added control for the pump via GPIO

This commit is contained in:
2021-01-27 03:00:55 +03:00
parent ca45f3e1dd
commit 36e3ebda77
2 changed files with 25 additions and 12 deletions

25
main.py
View File

@@ -3,20 +3,16 @@ import cv2
import numpy as np import numpy as np
import importlib.util import importlib.util
from threading import Thread from threading import Thread
import RPi.GPIO as GPIO
import time import time
from image import * from image import *
from pump import *
IM_WIDTH = 1280
IM_HEIGHT = 720
camera_type = 'usb'
MODEL_NAME = 'TFLite_model' MODEL_NAME = 'TFLite_model'
PATH_TO_CKPT = os.path.join(os.getcwd(), MODEL_NAME, 'detect.tflite') PATH_TO_CKPT = os.path.join(os.getcwd(), MODEL_NAME, 'detect.tflite')
PATH_TO_LABELS = os.path.join(os.getcwd(), MODEL_NAME, 'labelmap.txt') PATH_TO_LABELS = os.path.join(os.getcwd(), MODEL_NAME, 'labelmap.txt')
min_conf_threshold = 0.5 min_conf_threshold = 0.5
NUM_CLASSES = 90
# Import TensorFlow libraries # Import TensorFlow libraries
# If tflite_runtime is installed, import interpreter from tflite_runtime, else import from regular tensorflow # If tflite_runtime is installed, import interpreter from tflite_runtime, else import from regular tensorflow
@@ -89,7 +85,7 @@ class VideoStream:
input_mean = 127.5 input_mean = 127.5
input_std = 127.5 input_std = 127.5
def pet_detector(frame, detection_time): def pet_detector(frame, detection_time, cooldown):
frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
frame_resized = cv2.resize(frame_rgb, (width, height)) frame_resized = cv2.resize(frame_rgb, (width, height))
input_data = np.expand_dims(frame_resized, axis=0) input_data = np.expand_dims(frame_resized, axis=0)
@@ -107,16 +103,20 @@ def pet_detector(frame, detection_time):
for i in range(len(scores)): for i in range(len(scores)):
if ((scores[i] > min_conf_threshold) and (scores[i] <= 1.0)): if ((scores[i] > min_conf_threshold) and (scores[i] <= 1.0)):
obj_name = labels[int(classes[i])] obj_name = labels[int(classes[i])]
#print(obj_name)
if obj_name == 'cat' or obj_name == 'teddy bear': if obj_name == 'cat' or obj_name == 'teddy bear':
detection_time += 1 detection_time += 1
else: else:
detection_time = 0 detection_time = 0
no_cat() no_cat()
print(detection_time) if (time.process_time() - cooldown > 5):
if detection_time >= 2: switch_pump(False)
cat()
print('Looking at a cat for %s frames' % detection_time)
if detection_time >= 2:
cooldown = time.process_time()
cat()
switch_pump(True)
return frame, detection_time return frame, detection_time
@@ -124,9 +124,10 @@ videostream = VideoStream(resolution=(resW, resH), framerate=30).start()
time.sleep(1) time.sleep(1)
detection_time = 0 detection_time = 0
cooldown = 0
while(True): while(True):
frame = videostream.read() frame = videostream.read()
frame, detection_time = pet_detector(frame, detection_time) frame, detection_time = pet_detector(frame, detection_time, cooldown)
cv2.imshow('Cat-detector', frame) cv2.imshow('Cat-detector', frame)
if cv2.waitKey(1) == ord('q'): if cv2.waitKey(1) == ord('q'):

12
pump.py Normal file
View File

@@ -0,0 +1,12 @@
import RPi.GPIO as GPIO
import time
pin = 4
GPIO.setmode(GPIO.BCM)
GPIO.setup(pin, GPIO.OUT)
def switch_pump(state):
GPIO.output(pin, state)