v0.2

GiaoRobot.py

@@ -5,6 +5,7 @@
 from push import Push
 from storage import Storage
 from camera import Camera
+from utils import draw_bboxes
 
 log = Logger.getLogger("Giao")
 
@@ -34,17 +35,36 @@ def main():
                 if diff < config.MonitorDiffThreshold:
                     time.sleep(0.2)
                     continue
-                bboxes = camera.detectFaces(imgRegion)
-                log.info("bboxes len %d", len(bboxes))
-                if len(bboxes) > 0 and (time.time() - last_send_time) > 30:
+                faceNum, predType, predName = camera.detectFaces(imgRegion)
+                log.info("faceNum: %d", faceNum)
+                if faceNum > 0:
+                    if predType[1]:
+                        typeStr = '确定'
+                    else:
+                        typeStr = '可能'
+                    if predType[0] == 0:
+                        typeStr += '进入'
+                    elif predType[0] == 1:
+                        typeStr += '离开'
+                    else:
+                        typeStr += '站起来了'
+                    if predName[1]:
+                        nameStr = '确定是'
+                    else:
+                        nameStr = '可能是'
+                    nameStr += config.PersonNames[predName[0]]
+                    content = nameStr + '，' + typeStr
+                    log.info('push content %s', content)
+
+                if faceNum > 0 and (time.time() - last_send_time) > 30:
                     last_send_time = time.time()
                     saveName = str(time.strftime('images/%Y%m%d_%H_%M_%S.png'))
                     cv2.imwrite(saveName, img)
                     saveName = str(time.strftime('images/%Y%m%d_%H_%M_%S.jpg'))
                     cv2.imwrite(saveName, img)
                     url = storage.saveImage(saveName)
                     log.info("send giao! %s" % url)
-                    push.sendImage(config.PushTitle, url)
+                    push.sendImage(config.PushTitle, content, url)
         except Exception as e:
             log.error("error: %s", e)
 

camera.py

@@ -1,7 +1,9 @@
 from tinyface import TinyFace
+from recognition import Recognition
 import cv2, time
 import numpy as np
 from logger import Logger
+from utils import draw_bboxes
 
 log = Logger.getLogger('camera')
 
@@ -10,6 +12,7 @@ class Camera(object):
 
     def __init__(self, modelPath, device='cpu'):
         self.TF = TinyFace(modelPath, device=device)
+        self.REC = Recognition(device=device)
         self.lastImg = None
 
     def detectDiff(self, img):
@@ -36,4 +39,9 @@ def detectDiff(self, img):
     def detectFaces(self, img):
         imgRegion = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
         bboxes = self.TF.detect_faces(imgRegion, conf_th=0.9, scales=[1])
-        return bboxes
+        bboxesLength = len(bboxes)
+        predType, predName = (-1, True), (-1, True)
+        if bboxesLength> 0:
+            imgRegion = draw_bboxes(img, bboxes, thickness=1)
+            predType, predName = self.REC.detect(imgRegion, True)
+        return bboxesLength, predType, predName

config.py.example

@@ -3,29 +3,34 @@ dingTalkWebhookAccessToken = 'xxx'
 
 # the region size need to monitor
 # [x, y, w, h]
-MonitorRegion = [-1, -1, -1, -1]
+MonitorRegion = {'x': 270, 'y': 185, 'w': 220, 'h': 120}
+
+# diff threshold
+MonitorDiffThreshold = 300
 
 # TinyFace model path
 TinyFaceModelPath = './checkpoint_50.pth'
 # cpu or cuda
 TinyFaceModelDevice = 'cpu'
 
 # Sina SCS access key, secret key and bucket name
-SinaSCSAccessKey = 'xxx'
-SinaSCSSecretKey = 'xxx'
-SinaSCSBucketName = 'xxx'
-SinaSCSBucketUrl = 'http://xxx.xxx/xxx/%s'
+SinaSCSAccessKey = 'xx'
+SinaSCSSecretKey = 'xx'
+SinaSCSBucketName = 'xx'
+SinaSCSBucketUrl = 'http://abc.com/%s'
 
 # push title
 PushTitle = '[进入预警]'
 
+PersonNames = {0: 'xx', 1: 'xx', 2: 'xxx', 3: 'xxxx', 4: 'asd', 5: 'asd', 6: '2sda', 7: 'xa', 8: 'sad'}
+
 # Logger config json text
 LoggerJsonConfig = '''{
     "version":1,
     "disable_existing_loggers":false,
     "formatters":{
         "simple":{
-            "format":"[%(asctime)s][%(funcName)-15s:%(lineno)-4s] - %(name)-10s - %(levelname)-5s : %(message)s"
+            "format":"[%(asctime)s][%(funcName)-15s:%(lineno)-4s] - %(name)-15s - %(levelname)-5s : %(message)s"
         }
     },
     "handlers":{

push.py

@@ -7,8 +7,8 @@ def __init__(self, token):
         self.d = DingtalkChatbot(
             'https://oapi.dingtalk.com/robot/send?access_token=%s' % token)
 
-    def sendImage(self, title, url, is_at_all=False):
-        self.d.send_markdown(title=title, text='Giao, Giao, Giao!\n![Giao](' + url + ')\n', is_at_all=is_at_all)
+    def sendImage(self, title, content, url, is_at_all=False):
+        self.d.send_markdown(title=title, text=content + '\n![Giao](' + url + ')\n', is_at_all=is_at_all)
 
     def sendMessage(self, msg, is_at_all=False):
         self.d.send_text(msg=msg, is_at_all=is_at_all)

recognition.py

@@ -0,0 +1,54 @@
+import torch
+from torchvision.transforms import Compose, ToTensor, Normalize
+from torchvision import models
+import time, glob
+import cv2
+import numpy as np
+from tools.model import ResNet50
+import json
+from tools.getTrainDataSet import person
+from logger import Logger
+import config
+
+log = Logger.getLogger('Recognition')
+
+
+class Recognition(object):
+
+    def __init__(self, device='cpu'):
+        self.device = device
+        tstamp = time.time()
+        log.info('[Recognition] loading with %s', self.device)
+        resnet = models.resnet50()
+        resnet.load_state_dict(torch.load('tools/resnet50-19c8e357.pth'))
+        # self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+        net = ResNet50(resnet, 12)
+        net = net.to(self.device)
+        net.load_state_dict(torch.load('tools/models/latestModel.pth', map_location=torch.device(self.device)))
+        self.net = net
+        self.transform = Compose([ToTensor(), Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))])
+        log.info('[Recognition] finished loading (%.4f sec)' % (time.time() - tstamp))
+
+    def detect(self, img, bgr2rgb=True):
+        tstamp = time.time()
+        if bgr2rgb:
+            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+        img = self.transform(img).unsqueeze_(0)
+        img = img.to(self.device)
+        predLabel = self.net(img)
+        predictions = predLabel.data.cpu().numpy()[0]
+        predType = predictions[0:3]
+        predTypeId = np.where(predType > np.percentile(predType, 90))[0][0]
+        predNames = predictions[3:]
+        predNamesId = np.where(predNames > np.percentile(predNames, 90))[0][0]
+        log.info('predictions: ' + str(predictions))
+        ret1 = (predTypeId, True if predType[predTypeId] > 0 else False)
+        ret2 = (predNamesId, True if predNames[predNamesId] > 0 else False)
+        log.info('recognition using %.6f sec' % (time.time() - tstamp))
+        return ret1, ret2
+
+
+if __name__ == '__main__':
+    rec = Recognition()
+    # img = cv2.imread('tools/111.png')
+    # rec.detect(img)

tools/getTrainDataSet.py

@@ -1,10 +1,16 @@
 import cv2
 import tkinter as tk
 import json, glob
-import os
+import os, sys
+sys.path.append('..')
+import config
 
-person = {0: 'wk', 1: 'dln', 2: 'kml', 3: 'zsy', 4: 'hh', 5: 'wly', 6: 'lch', 7: 'wzh', 8: 'ys'}
 
+person = config.PersonNames
+
+index = 0
+
+l = [0] * 12
 
 def main():
     window = tk.Tk()
@@ -33,10 +39,6 @@ def main():
     if len(needHandleList) == 0:
         exit(0)
 
-    index = 0
-
-    l = [0] * 12
-
     img = cv2.imread(imagesWithBox + '/' + needHandleList[index])
     cv2.imshow('Image', img)
 
@@ -53,14 +55,14 @@ def onClick(data):
                 content.set(" " * 100)
                 content.set("Finished!")
                 return
-            if l[0] == 0 and l[1] == 0 and l[2] == 0:
-                return
+            # if l[0] == 0 and l[1] == 0 and l[2] == 0:
+            #     return
             mark = False
             for i in l[3:]:
                 if i == 1:
                     mark = True
-            if not mark:
-                return
+            # if not mark:
+            #     return
             labels[needHandleList[index].replace('.png', '')] = l
             with open(labelsJson, 'w') as file:
                 json.dump(labels, file)

tools/testModel.py

@@ -6,16 +6,24 @@
 import numpy as np
 from model import ResNet50
 import json
-from getTrainDataSet import person
+import os, sys
+
+sys.path.append('..')
+import config
+
+person = config.PersonNames
 
 OUT_DIM = 12
 
 resnet = models.resnet50()
 resnet.load_state_dict(torch.load('./resnet50-19c8e357.pth'))
 
+device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+
 net = ResNet50(resnet, OUT_DIM)
+net = net.to(device)
 
-net.load_state_dict(torch.load('models/latestModel.pth'))
+net.load_state_dict(torch.load('models/latestModel.pth', map_location=torch.device(device)))
 
 trainDataPath = './imagesWithBox'
 allTrainDataList = glob.glob(trainDataPath + '/*.png')
@@ -44,12 +52,21 @@ def getStr(data):
     imgO = cv2.imread(path)
     img = cv2.cvtColor(imgO, cv2.COLOR_BGR2RGB)
     img = transform(img).unsqueeze_(0)
+    img = img.to(device)
     t1 = time.time()
     predLabel = net(img)
     gtLable = labels[name]
     print('use %.4f' % (time.time() - t1))
     predictions = predLabel.data.cpu().numpy()[0]
-    idx_list = np.where(predictions > np.percentile(predictions, 90))[0]
+    predType = predictions[0:3]
+    predTypeId = np.where(predType > np.percentile(predType, 90))[0][0]
+    print(predType, predTypeId)
+    predNames = predictions[3:]
+    predNamesId = np.where(predNames > np.percentile(predNames, 90))[0][0]
+    print(predNames, predNamesId)
+    # idx_list = [predTypeId, predNamesId + 3]
+    idx_list = [predTypeId, predNamesId + 3]
+
     # print(gtLable)
     gtStr = 'GT '
     for i in range(len(gtLable)):
@@ -64,9 +81,9 @@ def getStr(data):
         predStr += ' '
     print(predStr)
     imgO = cv2.putText(imgO, gtStr, (10, 20), cv2.FONT_HERSHEY_SIMPLEX,
-                      0.7, (0, 0, 0), 1, cv2.LINE_AA)
+                       0.7, (0, 0, 0), 1, cv2.LINE_AA)
     imgO = cv2.putText(imgO, predStr, (10, 50), cv2.FONT_HERSHEY_SIMPLEX,
                        0.7, (0, 0, 0), 1, cv2.LINE_AA)
     cv2.imshow("image", imgO)
     cv2.waitKey(0)
-
+    # break

tools/trainModel.py

@@ -76,6 +76,8 @@ def __getitem__(self, index):
 
 device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 
+net = net.to(device)
+
 trainDataPath = './imagesWithBox'
 allTrainDataList = glob.glob(trainDataPath + '/*.png')