main.py

import os
import cv2
import numpy as np
import datetime as dt
import sqlite3
import time
import requests
from dotenv import load_dotenv

from draw_poly import *
from utils import *
from recognizer import *

conf_threshold = 0

load_dotenv()
conn = sqlite3.connect('db.sqlite')
cur = conn.cursor()

backend = os.environ['BACKEND_URL']
yard = os.environ['YARD']
offline_mode = os.environ.get('OFFLINE', 'false').lower() == 'true'
if backend[-1] != '/':
    backend += '/'

def sync_db():
    if offline_mode:
        print("Offline mode. DB sync won't be perfomed.")
        return
    print('Refreshing DB...')
    resp = requests.get(backend+'api/yard/auto-numbers/?yard_id='+yard)
    if resp.status_code//100!=2:
        print('Cannot refresh DB. Retry in 1 min')
        tasks['sync_db'][2]=time.time()-tasks['sync_db'][1]+60
        return
    for auto_number in resp.json()['auto_numbers']:
        cur.execute("INSERT OR REPLACE INTO cars(text, inside) VALUES (?,?)",
                    (auto_number, 0))
    conn.commit()
    print('DB has been refreshed successfully. See you in 5 min')

def remove_ignores():
    global ignore
    ignore = {k:v for k,v in ignore.items() if dt.datetime.now() - v[0] > dt.timedelta(minutes=1)}

ignore = {}
def process_detection(cap_number, text, rect):
    if text == ValueError:
        return
    if text in ignore:
        if ignore[text][1]!=cap_number:
##            cur.execute("INSERT OR REPLACE INTO cars(text) VALUES (?)", (text,))
            cur.execute("UPDATE cars SET inside=? WHERE text=?",
                        (1-cap_number, text))
            conn.commit()
        if dt.datetime.now() - ignore[text][0] > dt.timedelta(minutes=1):
            del ignore[text]
        else:
            return
    print(text)
    query = cur.execute("SELECT text FROM cars WHERE text = ?", (text,))
    if query.fetchone() is not None:
        if cap_number==exit_id:
            return
        cur.execute("INSERT OR REPLACE INTO cars(text) VALUES (?)", (text,))
        cur.execute("UPDATE cars SET inside=1-inside WHERE text=?", (text,))
        conn.commit()
        if offline_mode:
            print("Offline mode. Cannot decide anything.")
        else:
            resp = requests.post(backend+'api/yard/history/', json=
                          {
                              'auto_number': text,
                              'event_type': ('entry','exit')[cap_number],
                              'yard_id': os.environ['YARD']
                           })
            if resp.status_code//100==2:
                if 'error' not in resp.json():
                    print('Passed', text)
                else:
                    print('Go away', text)
            else:
                print('Error', resp.status_code)
                print(resp.text,'\n')
    ignore[text] = (dt.datetime.now(), cap_number)

# keep in mind: cam 0 must watch the entry; cam 1 watch the exit
entry_id = input('Entry camera id: ').strip()
exit_id = input('Exit camera id: ').strip()
caps = (cv2.VideoCapture(int(entry_id) if entry_id else 0),
        cv2.VideoCapture(int(exit_id) if exit_id else 1))

##caps = (cv2.VideoCapture("C:\\Users\\USER\\Downloads\\IMG_2907_2fps.mp4"),
##       cv2.VideoCapture(1))
if not all(cap.isOpened() for cap in caps):
    print("Could't connect to camera")
    exit()

print('Captured both cameras. Press Q to close all windows')
draw_poly = 1
polygons, zones = zip(load_clipping_polygon('polygons.json','entry'),
                      load_clipping_polygon('polygons.json','exit'))
polygon_lists = [[list(map(int,i)) for i in polygon.exterior.coords]
                for polygon in polygons]
recognizer = Recognizer()

start_pos = (10, 10)
end_pos = (start_pos[0]+5, start_pos[1]+5)
height, width = 480, 600
frame = np.zeros((height, width, 3), np.uint8)
tasks = {
    'sync_db': [sync_db, 300, 0],
    'remove_ignores': [remove_ignores, 100, 0]
}
while 1:
##    clock = time.time()
    for cap_number, cap in enumerate(caps):
        frame = np.zeros((height, width, 3), np.uint8)
        height, width, _ = frame.shape
        ret, frame = cap.read()
        if not ret:
            cv2.rectangle(frame, start_pos, end_pos, (0,0,255), -1)
            continue
        for text, rect, conf in recognizer(frame):
            if conf<conf_threshold:
                continue
            if rectangle_intersects_clipping_zone(rect, polygons[cap_number]):
                cv2.rectangle(frame, *rect, (0,0,255), 3)
                continue
            text = correct_text(text)
            if text is ValueError:
                cv2.rectangle(frame, *rect, (150,150,150), 1)
                continue
            process_detection(cap_number, text, rect)
            cv2.rectangle(frame, *rect, 
                        (255,255,255) if text in ignore else (255,0,0), 3)
            cv2.putText(frame, f'{text} {round(conf,2)}',
                        (rect[0][0],rect[0][1]-12), 2, 1,(255,255,255))
        if draw_poly:
            for i in range(len(polygon_lists[cap_number])-1):
                cv2.line(frame,
                                 polygon_lists[cap_number][i],
                                 polygon_lists[cap_number][i+1],
                                 (0,0,255), 2)
        frame = cv2.resize(frame, (width, height), interpolation=cv2.INTER_AREA)
        cv2.rectangle(frame, start_pos, end_pos, (0,255,0), -1)
        cv2.imshow(f'Camera {cap_number}. Cropping {zones[cap_number]} the line', frame)

##    print(round(1/(time.time()-clock), 2), 'fps')
##    clock = time.time()

    for task in tasks.values():
        if time.time()-task[2]>task[1]:
            task[2]=time.time()
            task[0](**(task[3] if len(task)>3 else {}))

    if cv2.waitKey(1) & 0xFF == ord('q'):
        for i in range(len(caps)):
            caps[i].release()
        cv2.destroyAllWindows()
        break
cur.close()