Ejemplo n.º 1
0
class Detector:
    def __init__(self, camera):
        self.initModel()
        self.initTwitter()
        self.frameTime = 0
        self.frames = 0
        self.lastImageTime = 0
        self.fps = 30
        self.lastTweetTime = 0
        self.vr = VideoReader(camera, self.fps)

    def initModel(self):
        self.interpreter = tf.lite.Interpreter(model_path=MODEL_FILE)
        self.interpreter.allocate_tensors()
        self.input_details = self.interpreter.get_input_details()
        self.output_details = self.interpreter.get_output_details()

        self.model_height = self.input_details[0]['shape'][1]
        self.model_width = self.input_details[0]['shape'][2]

        self.floating_model = (self.input_details[0]['dtype'] == np.float32)

    def initTwitter(self):
        with open(TWITTER_CONFIG) as f:
            data = json.load(f)
        auth = tweepy.OAuthHandler(data['apiKey'], data['apiSecret'])
        auth.set_access_token(data['token'], data['tokenSecret'])
        self.twitter = tweepy.API(auth)

    def detect(self, image):
        image = image.copy()

        # Test model on random input data.
        input_shape = self.input_details[0]['shape']
        image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        image_resized = cv2.resize(image_rgb,
                                   (self.model_width, self.model_height))
        input_data = np.expand_dims(image_resized, axis=0)
        if self.floating_model:
            input_data = (np.float32(input_data) - INPUT_MEAN) / INPUT_STD
        self.interpreter.set_tensor(self.input_details[0]['index'], input_data)

        # Run the network
        self.interpreter.invoke()

        boxes = self.interpreter.get_tensor(self.output_details[0]['index'])[0]
        classes = self.interpreter.get_tensor(
            self.output_details[1]['index'])[0]
        scores = self.interpreter.get_tensor(
            self.output_details[2]['index'])[0]

        return boxes, scores, classes

    def triggerDetection(self, frame, orig_frame):
        self.log('Norppa detected!')

        self.lastDetectTime = time.time()

        # Only react every 30 seconds
        if (time.time() - self.lastImageTime) < 30:
            return

        # Save the original image and frame with detection boxes
        filename = "images/frame%s.jpg" % datetime.now().strftime(
            "%Y%m%d-%H%M%S")
        o_filename = "images/orig_frame%s.jpg" % datetime.now().strftime(
            "%Y%m%d-%H%M%S")
        cv2.imwrite(filename, frame)
        cv2.imwrite(o_filename, orig_frame)
        self.log("Wrote %s" % filename)

        # Tweet
        self.tweetNorppaIsLive(frame)

        self.lastImageTime = time.time()

    def tweetNorppaIsLive(self, frame):
        # only tweet every TWEET_TIME_MINUTES minutes
        if time.time() - self.lastTweetTime < TWEET_TIME_MINUTES * 60:
            return

        # Get random message from messages.json
        with open(MESSAGES_FILE, 'r') as f:
            messages = json.load(f)

        randomMessage = random.choices(messages)[0].replace('#', 'hashtag-')

        filename = 'uploaded_image.jpg'
        cv2.imwrite(filename, frame)
        self.twitter.update_with_media(filename, status=randomMessage)
        self.log('Tweeted: %s' % randomMessage)
        os.remove(filename)
        self.lastTweetTime = time.time()

    def checkCoordinates(self, x, y, x2, y2):
        # You can add coordinate & width/height checks here
        width = x2 - x
        height = y2 - y
        area = width * height
        return area < 350000 and area > 10000

    def log(self, msg):
        date = datetime.now().strftime("%Y-%m-%d-%H:%M:%S")
        print(date, msg)

    def run(self):
        print("Ready")
        self.vr.start()
        self.startTime = time.time()
        frameCounter = 0
        lastFrameTime = 0
        detected = False
        detections = 0

        # Wait for the stream to start
        while not self.vr.more() or time.time() - self.startTime > 10:
            time.sleep(0.5)

        frame = None
        while True:
            if not self.vr.more():
                print('asd')
                break

            # Read the latest frame
            frame = np.copy(self.vr.read())
            frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            frameCounter += 1

            lastFrameTime = time.time()

            orig_frame = frame.copy()
            detected = False

            # Get detection boxes and scores from tensorflow
            boxes, scores, classes = self.detect(frame)

            # Debug logging
            if scores[0] > 0.4:
                self.log(scores[0])

            # Go through the boxes
            for i in range(len(scores)):
                score = scores[i]
                box = boxes[i]

                y = int(box[0] * HEIGHT)
                x = int(box[1] * WIDTH)
                y2 = int(box[2] * HEIGHT)
                x2 = int(box[3] * WIDTH)

                if score < THRESHOLD:
                    continue

                # Some sanity checks
                if not self.checkCoordinates(x, y, x2, y2):
                    continue

                # Norppa detected here!
                detected = True
                self.log(score)
                print((x, y), (x2, y2))

                # Draw the detection box to the frame
                cv2.rectangle(frame, (x, y), (x2, y2), (0, 255, 0), 2)
                label = '%s: %d%%' % ('norppa', int(score * 100))
                labelSize, baseLine = cv2.getTextSize(label,
                                                      cv2.FONT_HERSHEY_SIMPLEX,
                                                      0.7, 2)
                label_ymin = max(y, labelSize[1] + 10)
                # Draw white box to put label text in
                cv2.rectangle(frame, (x, label_ymin - labelSize[1] - 10),
                              (x + labelSize[0], label_ymin + baseLine - 10),
                              (255, 255, 255), cv2.FILLED)
                cv2.putText(frame, label, (x, label_ymin - 7),
                            cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 0), 2)
            if detected:
                detections += 1
            else:
                detections = 0

            # Only detect after the second frame in a row that breaks the threshold
            if detections >= 2:
                self.triggerDetection(frame, orig_frame)

            if frameCounter % 1000 == 0:
                self.log('1000 frames')
                frameCounter = 0
            frameTime = time.time() - lastFrameTime
            if frameTime < 1 / self.fps:
                time.sleep(1 / self.fps - frameTime)
        self.vr.stop()
        time.sleep(1)
Ejemplo n.º 2
0
def main(video, json_logger=None, visualise_callback=False, mode=None, inferencer=None, model_name=None, image_size=None, model_alpha=None, shot_name=None, preset_name=None, prob_threshold=None, iou_threshold=None, union_threshold=None, model_precision=None, inference_device=None, num_shots_hint=None, video_descriptor=None, gender_estimation=False, emotions_recognision=False, run_faceid=False, create_feauture_list=None, faceid_search=False):
    if faceid_search:
        run_faceid = True
    if not(create_feauture_list is None):
        feauture_list = []
        run_faceid = True
        if preset_name is None:
            preset_name = 'accuracy'

    if not(mode is None):
        try:
            mode = str(mode)
        except (TypeError, ValueError):
            raise ErrorSignal(invalid_argument_value)
    else:
        mode = default_mode

    if not(json_logger is None):
        try:
            json_logger = str(json_logger)
        except (TypeError, ValueError):
            raise ErrorSignal(invalid_argument_value)

    exit_code = 1

    reader = VideoReader(video, width_hint, height_hint)
    detector = FaceDetection(reader.get_width(), reader.get_height(), inferencer=inferencer, model_name=model_name, image_size=image_size, model_alpha=model_alpha, shot_name=shot_name, preset_name=preset_name, prob_threshold=prob_threshold, iou_threshold=iou_threshold, union_threshold=union_threshold, model_precision=model_precision, inference_device=inference_device, num_shots_hint=num_shots_hint)
    if gender_estimation:
        gender_estimatior = GenderEstimation(reader.get_width(), reader.get_height(), inferencer=inferencer)
    if run_faceid:
        faceid = FaceID(reader.get_width(), reader.get_height(), inferencer=inferencer)
    if emotions_recognision:
        emotions_recogniser = EmotionsRecognision(reader.get_width(), reader.get_height(), inferencer=inferencer)
    if faceid_search:
        try:
            faceid_dict_f = open(faceid_dict, "r")
        except IOError:
            raise ErrorSignal(faceid_dict_missing)
        faceid_json = json.load(faceid_dict_f)
        faceid_dict_f.close()
        faceids = []
        for name in faceid_json:
            faceids.append((name, faceid_json[name]['threshold'], np.array(faceid_json[name]['feautures'])))

    callbacks = []

    if visualise_callback:
        callbacks.append(Visualiser(reader.get_width(), reader.get_height()))

    if not(json_logger is None):
        callbacks.append(JSONLogger(reader.get_width(), reader.get_height(), json_logger))

    if len(callbacks) == 0:
        reader.release()
        raise ErrorSignal(nothing_to_do)

    if mode == 'byframe':
        frame_number = 0
        try:
            while True:
                frame = reader.read()
                boxes = detector.detect(frame)
                genders = None
                faceid_feautures = None
                emotions = None
                names = None
                if gender_estimation:
                    genders = gender_estimatior.estimate(frame, boxes)
                if emotions_recognision:
                    emotions = emotions_recogniser.recognise(frame, boxes)
                if run_faceid:
                    faceid_feautures = faceid.feautures(frame, boxes)
                if faceid_search:
                    names = []
                    for i in range(faceid_feautures.shape[0]):
                        face_names = []
                        for fid in faceids:
                            d = np.min(cosine_dist_norm(faceid_feautures[i].reshape(1, 128), fid[2]))
                            if d < fid[1]:
                                face_names.append(fid[0])
                        names.append(face_names)
                if not(create_feauture_list is None):
                    if faceid_feautures.shape[0] > 1:
                        raise ErrorSignal(only_one_face_required_in_create_feauture_list_mode)
                    elif faceid_feautures.shape[0] == 1:
                        feauture_list.append(faceid_feautures[0])
                frame_number += 1
                for callback in callbacks:
                    try:
                        callback.call(video=video, frame=frame, boxes=boxes, frame_number=frame_number, genders=genders, emotions=emotions, faceid_feautures=faceid_feautures, names=names)
                    except QuitSignal as ret:
                        exit_code = int(ret)
                        break
                else:
                    continue
                break
        except KeyboardInterrupt:
            exit_code = 0
            pass
        reader.release()
        for callback in callbacks:
            callback.destroy()
    elif mode == 'realtime':
        raise NotImplementedError()
    else:
        raise ErrorSignal(unknown_mode)

    if not(create_feauture_list is None):
        try:
            create_feauture_list = str(create_feauture_list)
        except (KeyError, ValueError):
            raise ErrorSignal(invalid_argument_value)
        feautures, threshold = feauture_select(feauture_list)
        try:
            faceid_dict_f = open(faceid_dict, "r")
        except IOError:
            raise ErrorSignal(faceid_dict_missing)
        faceid_json = json.load(faceid_dict_f)
        faceid_dict_f.close()
        faceid_json[create_feauture_list] = {
            'threshold': float(threshold),
            'feautures': feautures.tolist()
        }
        try:
            faceid_dict_f = open(faceid_dict, "w")
        except IOError:
            raise ErrorSignal(faceid_dict_missing)
        json.dump(faceid_json, faceid_dict_f)
        faceid_dict_f.close()
        print(faceid_success_prefix + create_feauture_list + faceid_success_suffix)

    return exit_code