forked from mharizanov/AI-ipcam
-
Notifications
You must be signed in to change notification settings - Fork 0
/
ai-ipcam.py
107 lines (91 loc) · 4 KB
/
ai-ipcam.py
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
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
#!/usr/bin/env python
from darkflow.net.build import TFNet
import cv2
import os
import random
from subprocess import Popen
import time
from PIL import Image,ImageDraw
import numpy as np
import paho.mqtt.client as mqtt
import argparse
#RTSP captured frame frame_filename; preferably on RAM drive
#diskutil erasevolume HFS+ 'RAM Disk' `hdiutil attach -nomount ram://20480` on a MAC
frame_filename = '/Volumes/RAM Disk/frame'+str(random.randint(1,99999))+'.jpeg'
#threshold parameter is below, keeping it too low will result in recognition errors
options = {"model": "cfg/yolo.cfg", "load": "/Users/mharizanov/Desktop/cv/darknet/yolo.weights", "threshold": 0.45}
parser=argparse.ArgumentParser()
parser.add_argument(
"--watch", # name on the parser - drop the `--` for positional/required parameters
nargs="*", # 0 or more values expected => creates a list
type=str,
default=['person', 'cat', 'dog', 'bird'], # default if nothing is provided
)
parser.add_argument(
"--stream", # name on the parser - drop the `--` for positional/required parameters
type=str,
default= 'rtsp://192.168.1.153:554/onvif1', # default if nothing is provided
)
parser.add_argument(
"--broker", # name on the parser - drop the `--` for positional/required parameters
type=str,
default= 'broker.hivemq.com', # default if nothing is provided
)
parser.add_argument(
"--topic", # name on the parser - drop the `--` for positional/required parameters
type=str,
default= 'harizanov/home/camera1/', # default if nothing is provided
)
parser.add_argument(
"--showimage", # name on the parser - drop the `--` for positional/required parameters
type=str,
default='no', # default if nothing is provided
)
# parse the command line
args = parser.parse_args()
watch_list=args.watch
rtsp_stream=args.stream
broker_address=args.broker
mqtt_topic=args.topic
showimageflag=args.showimage
print("Watching for: %r" % watch_list)
print("Stream: %r" % rtsp_stream)
print("MQTT broker: %r" % broker_address)
print("MQTT topic: %r" % mqtt_topic)
print("Show image: %r" % showimageflag)
client = mqtt.Client("cameraclient_"+str(random.randint(1,99999999))) #create new instance
client.connect(broker_address)
#start a ffmpeg process that captures one frame every 5 seconds
p = Popen(['ffmpeg', '-loglevel', 'panic', '-rtsp_transport', 'udp', '-i', rtsp_stream, '-f' ,'image2' ,'-pix_fmt', 'yuvj420p', '-r', '1/5' ,'-updatefirst', '1', frame_filename])
tfnet = TFNet(options)
while True:
try:
curr_img = Image.open( frame_filename )
curr_img_cv2 = cv2.cvtColor(np.array(curr_img), cv2.COLOR_RGB2BGR) #is the frame good and can be opened?
os.remove(frame_filename) #delete frame once it is processed, so we don't reprocess the same frame over and over
except: # ..frame not ready, just snooze for a bit
time.sleep(1)
continue
result = tfnet.return_predict(curr_img_cv2)
print(result)
client.publish(mqtt_topic, "".join([str(x) for x in result]) ) #publish
saveflag=False
namestr=''
draw = ImageDraw.Draw(curr_img)
for det_object in result:
if any(det_object['label'] in s for s in watch_list):
draw.rectangle([det_object['topleft']['x'], det_object['topleft']['y'],det_object['bottomright']['x'], det_object['bottomright']['y']], outline=(255, 255, 0))
draw.text([det_object['topleft']['x'], det_object['topleft']['y'] - 13], det_object['label']+' - ' + str( "{0:.0f}%".format(det_object['confidence'] * 100) ) , fill=(255, 255, 0))
saveflag=True
namestr+='_'+str(det_object['label'])
if saveflag == True:
curr_img.save('images/'+str(int(time.time()))+namestr+'.jpg')
saveflag=False
if showimageflag!='no':
curr_img_cv2=cv2.cvtColor(np.array(curr_img), cv2.COLOR_RGB2BGR)
cv2.imshow("Security Feed", curr_img_cv2)
if cv2.waitKey(1) & 0xFF == ord('q'): # wait for 1 millisecond
break
continue
time.sleep(1)
p.terminate()