video_capture.py

import cv2
from datetime import datetime, timedelta
import numpy as np
import keras
from utils.preprocessing import fix_singular_shape, norm_by_imagenet
import json
import time

from cameras import restaurants
from db import models
from db.database import SessionLocal
from routers.websocket import manager


async def handle_cameras():
    model = keras.models.load_model('assets', compile=False)
    db = SessionLocal()
    for restaurant in restaurants:
        for camera in restaurant["cameras"]:
            mask = np.zeros((720, 1280, 3), dtype=np.float32)
            cv2.fillPoly(mask, np.array(camera["mask_points"]), (255, 255, 255))
            camera["mask"] = mask

    while True:

        start_timestamp = time.time()
        current_date = datetime.fromtimestamp(start_timestamp)
        weekday, current_time = current_date.weekday(), current_date.time()

        for restaurant in restaurants:

            is_open = db.query(
                models.OpeningHours).filter(
                models.OpeningHours.place == restaurant["restaurant"],
                models.OpeningHours.day == weekday,
                models.OpeningHours.open_time <= current_time,
                models.OpeningHours.close_time >= current_time).first() is not None

            if is_open:
                count_prediction = 0
                open_checkouts = 0
                cams_working = True

                for camera in restaurant["cameras"]:
                    cap = cv2.VideoCapture(f'rtsp://{camera["user"]}:{camera["password"]}@{camera["IP"]}/{camera["stream"]}')
                    if cams_working and cap.isOpened():
                        _, frame = cap.read()
                        masked_img = cv2.bitwise_and(
                            frame.astype(np.float32), camera["mask"])
                        treated_img = fix_singular_shape(masked_img, 16)
                        input_image = np.expand_dims(
                            np.squeeze(
                                norm_by_imagenet(
                                    np.array(
                                        [treated_img]))),
                            axis=0)
                        pre_pred = time.time()
                        pred_map = np.squeeze(model.predict(input_image, verbose=0))
                        print(time.time() - pre_pred)
                        count_prediction += np.sum(pred_map)
                        for caisse in camera["caisses"]:
                            if np.sum(pred_map[caisse["x1"] // 2:caisse["x2"] // 2, caisse["y1"] // 2:caisse["y2"] // 2]) > 0.5:
                                open_checkouts += 1
                    else:
                        cams_working = False
                    cap.release()

                if cams_working:
                    waiting_time = timedelta(
                        seconds=restaurant['b_factor'] +
                        int(count_prediction *
                            restaurant['a_factor'] / max(open_checkouts, 1)))
                    db_record = models.Records(
                        place=restaurant['restaurant'],
                        date=current_date,
                        density=int(count_prediction),
                        waiting_time=waiting_time)
                    db.add(db_record)
                    db.commit()
                    await manager.broadcast(json.dumps({"type": "data"}))
        time.sleep(60 - time.time() + start_timestamp)