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
from db import models

from db.database import SessionLocal

db = SessionLocal()
model = keras.models.load_model('assets')

frame_gap = 450

cameras = [{
    "place": "local",
    "IP": "10.148.38.9",
    "user": "viarezocam",
    "password": "superponey",
    "stream": "stream1",
    "framegap": 900, # 60 * camera frequency
    "count": 0, # mandatory
    "cap": None
}]

for camera in cameras:
    camera.cap = cv2.VideoCapture(f"rtsp://{camera.user}:{camera.password}@{camera.ip}/{camera.stream}")

while True:
    for camera in cameras:
        if camera.cap.isOpened():
            ret, frame = camera.cap.read()
            if ret and camera.count % camera.frame_gap == 0:
                current_time = datetime.now()
                treated_img = fix_singular_shape(frame, 16)
                input_image = np.expand_dims(np.squeeze(norm_by_imagenet(np.array([treated_img]))), axis=0)
                pred_map = np.squeeze(model.predict(input_image))
                count_prediction = np.sum(pred_map)
                waiting_time = timedelta(seconds=120 + count_prediction * 30)
                record = {"place": camera.place,
                        "date": current_time,
                        "density": count_prediction,
                        "waiting_time": waiting_time}
                db_record = models.Records(**record)
                db.add(db_record)
                db.commit()
                db.refresh(db_record)
            camera.count += 1
        else:
            camera.cap.release()