Skip to content

Commit 6568ef2

Browse files
committed
Added Cat Detector Code and STL
1 parent d5aeaa3 commit 6568ef2

File tree

2 files changed

+177
-0
lines changed

2 files changed

+177
-0
lines changed

Lobe_Cat_Detector/Cat Detector.stl

87.7 KB
Binary file not shown.
+177
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,177 @@
1+
import time
2+
from enum import Enum, auto
3+
import board
4+
from digitalio import DigitalInOut, Direction, Pull
5+
import picamera
6+
import io
7+
from PIL import Image
8+
from lobe import ImageModel
9+
import os
10+
import adafruit_dotstar
11+
from datetime import datetime
12+
import pwmio
13+
from adafruit_motor import servo
14+
15+
LABEL_CAT = "Cat"
16+
LABEL_MULTI_CAT = "Multiple Cats"
17+
LABEL_NOTHING = "No Cats"
18+
SERVO_PIN = board.D12
19+
WARNING_COUNT = 3
20+
21+
pwm = pwmio.PWMOut(SERVO_PIN, duty_cycle=0, frequency=50)
22+
servo = servo.Servo(pwm, min_pulse=400, max_pulse=2400)
23+
24+
# Boiler Plate code for buttons and joystick on the braincraft
25+
BUTTON_PIN = board.D17
26+
JOYDOWN_PIN = board.D27
27+
JOYLEFT_PIN = board.D22
28+
JOYUP_PIN = board.D23
29+
JOYRIGHT_PIN = board.D24
30+
JOYSELECT_PIN = board.D16
31+
32+
buttons = [BUTTON_PIN, JOYUP_PIN, JOYDOWN_PIN,
33+
JOYLEFT_PIN, JOYRIGHT_PIN, JOYSELECT_PIN]
34+
35+
for i, pin in enumerate(buttons):
36+
buttons[i] = DigitalInOut(pin)
37+
buttons[i].direction = Direction.INPUT
38+
buttons[i].pull = Pull.UP
39+
button, joyup, joydown, joyleft, joyright, joyselect = buttons
40+
41+
42+
class Input(Enum):
43+
BUTTON = auto()
44+
UP = auto()
45+
DOWN = auto()
46+
LEFT = auto()
47+
RIGHT = auto()
48+
SELECT = auto()
49+
50+
51+
def get_inputs():
52+
inputs = []
53+
if not button.value:
54+
inputs.append(Input.BUTTON)
55+
if not joyup.value:
56+
inputs.append(Input.UP)
57+
if not joydown.value:
58+
inputs.append(Input.DOWN)
59+
if not joyleft.value:
60+
inputs.append(Input.LEFT)
61+
if not joyright.value:
62+
inputs.append(Input.RIGHT)
63+
if not joyselect.value:
64+
inputs.append(Input.SELECT)
65+
return inputs
66+
67+
DOTSTAR_DATA = board.D5
68+
DOTSTAR_CLOCK = board.D6
69+
70+
RED = (0, 0, 255)
71+
GREEN = (255, 0, 0)
72+
OFF = (0, 0, 0)
73+
74+
dots = adafruit_dotstar.DotStar(DOTSTAR_CLOCK, DOTSTAR_DATA, 3, brightness=0.1)
75+
76+
jingle_count = 0
77+
78+
def color_fill(color, wait):
79+
dots.fill(color)
80+
dots.show()
81+
time.sleep(wait)
82+
83+
def jingle_keys(jingle_hard=False):
84+
global jingle_count
85+
jingle_count += 1
86+
if jingle_count > WARNING_COUNT:
87+
jingle_hard = True
88+
delay = 0.5 if jingle_hard else 2
89+
loop = 5 if jingle_hard else 1
90+
travel = 180 if jingle_hard else 135
91+
for _ in range(0, loop):
92+
for angle in (0, travel):
93+
servo.angle = angle
94+
time.sleep(delay)
95+
servo.angle = None
96+
97+
def main():
98+
global jingle_count
99+
model = ImageModel.load('~/model')
100+
101+
# Check if there is a folder to keep the retraining data, if it there isn't make it
102+
if (not os.path.exists('./retraining_data')):
103+
os.mkdir('./retraining_data')
104+
105+
with picamera.PiCamera(resolution=(224, 224), framerate=30) as camera:
106+
stream = io.BytesIO()
107+
camera.start_preview()
108+
# Camera warm-up time
109+
time.sleep(2)
110+
label = ''
111+
while True:
112+
stream.seek(0)
113+
camera.annotate_text = None
114+
camera.capture(stream, format='jpeg')
115+
camera.annotate_text = label
116+
img = Image.open(stream)
117+
result = model.predict(img)
118+
label = result.prediction
119+
confidence = result.labels[0][1]
120+
camera.annotate_text = label
121+
print(f'\rLabel: {label} | Confidence: {confidence*100: .2f}%', end='', flush=True)
122+
123+
# Check if the current label is package and that the label has changed since last tine the code ran
124+
if label == LABEL_CAT:
125+
# Make Servo Jingle Keys
126+
jingle_keys()
127+
elif label == LABEL_MULTI_CAT:
128+
jingle_keys(True)
129+
elif label == LABEL_NOTHING:
130+
jingle_count = 0
131+
132+
time.sleep(0.5)
133+
134+
inputs = get_inputs()
135+
# Check if the joystick is pushed up
136+
if (Input.UP in inputs):
137+
color_fill(GREEN, 0)
138+
# Check if there is a folder to keep the retraining data, if it there isn't make it
139+
if (not os.path.exists(f'./retraining_data/{label}')):
140+
os.mkdir(f'./retraining_data/{label}')
141+
# Remove the text annotation
142+
camera.annotate_text = None
143+
144+
# File name
145+
name = datetime.now()
146+
# Save the current frame
147+
camera.capture(
148+
os.path.join(
149+
f'./retraining_data/{label}',
150+
f'{datetime.now().strftime("%Y-%m-%d_%H:%M:%S")}.jpg'
151+
)
152+
)
153+
154+
color_fill(OFF, 0)
155+
156+
# Check if the joystick is pushed down
157+
elif (Input.DOWN in inputs or Input.BUTTON in inputs):
158+
color_fill(RED, 0)
159+
# Remove the text annotation
160+
camera.annotate_text = None
161+
# Save the current frame to the top level retraining directory
162+
camera.capture(
163+
os.path.join(
164+
f'./retraining_data',
165+
f'{datetime.now().strftime("%Y-%m-%d_%H:%M:%S")}.jpg'
166+
)
167+
)
168+
color_fill(OFF, 0)
169+
170+
171+
if __name__ == '__main__':
172+
try:
173+
print(f"Predictions starting, to stop press \"CTRL+C\"")
174+
main()
175+
except KeyboardInterrupt:
176+
print("")
177+
print(f"Caught interrupt, exiting...")

0 commit comments

Comments
 (0)