126 lines
4.8 KiB
Python
126 lines
4.8 KiB
Python
import abc
|
|
|
|
import pandas as pd
|
|
from pathlib import Path
|
|
|
|
from dodgeballtracker.metadata import VideoMetadata
|
|
from dodgeballtracker.typing import TrackedDetectedObject, VideoTimestamp
|
|
from sqlite3 import Connection
|
|
|
|
import pandera as pa
|
|
|
|
|
|
class DataDumper(abc.ABC):
|
|
|
|
@abc.abstractmethod
|
|
def write_boxes(self, obj: list[TrackedDetectedObject], ts: VideoTimestamp,
|
|
frame_index: int) -> None:
|
|
pass
|
|
|
|
@abc.abstractmethod
|
|
def start(self, video_path: Path, hexa_hash: str) -> None:
|
|
pass
|
|
|
|
|
|
class SQLiteDataDumper(DataDumper):
|
|
def __init__(self, conn: Connection):
|
|
self._conn = conn
|
|
self._video_id = None
|
|
|
|
def _create_tables(self):
|
|
cur = self._conn.cursor()
|
|
|
|
cur.executescript('''
|
|
CREATE TABLE IF NOT EXISTS video (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
file_name TEXT NOT NULL,
|
|
file_sha256_hex TEXT NOT NULL,
|
|
count INTEGER NOT NULL DEFAULT 1,
|
|
|
|
UNIQUE(file_sha256_hex)
|
|
);
|
|
''')
|
|
|
|
cur.executescript('''
|
|
CREATE TABLE IF NOT EXISTS video_frame(
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
timestamp_ DOUBLE PRECISION NOT NULL,
|
|
index_ INTEGER NOT NULL,
|
|
video_id INTEGER NOT NULL,
|
|
nbr_detected_object INTEGER NOT NULL,
|
|
|
|
FOREIGN KEY(video_id) REFERENCES video(id),
|
|
UNIQUE(video_id, index_)
|
|
);
|
|
''')
|
|
|
|
cur.execute('''
|
|
CREATE TABLE IF NOT EXISTS tracked_detected_object(
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
xmin DOUBLE PRECISION NOT NULL,
|
|
ymin DOUBLE PRECISION NOT NULL,
|
|
xmax DOUBLE PRECISION NOT NULL,
|
|
ymax DOUBLE PRECISION NOT NULL,
|
|
|
|
bird_view_x DOUBLE PRECISION,
|
|
bird_view_y DOUBLE PRECISION,
|
|
|
|
score DOUBLE PRECISION NOT NULL,
|
|
category INTEGER NOT NULL,
|
|
object_id INTEGER NOT NULL,
|
|
video_frame_id INTEGER NOT NULL,
|
|
FOREIGN KEY(video_frame_id) REFERENCES video_frame(id)
|
|
);
|
|
''')
|
|
|
|
self._conn.commit()
|
|
|
|
def start(self,
|
|
video_path: Path,
|
|
hexa_hash: str,
|
|
metadata: VideoMetadata|None=None) -> None:
|
|
self._create_tables()
|
|
|
|
cur = self._conn.cursor()
|
|
res = cur.execute("""
|
|
INSERT INTO video(file_name, file_sha256_hex)
|
|
VALUES (?, ?)
|
|
ON CONFLICT(file_sha256_hex) DO UPDATE SET count = count + 1
|
|
RETURNING id
|
|
""", (video_path.name, hexa_hash))
|
|
self._video_id = res.fetchone()[0]
|
|
self._conn.commit()
|
|
|
|
def _is_empty_df(self, df):
|
|
return not df.shape[0]
|
|
|
|
@pa.check_io(df=TrackedDetectedObject)
|
|
def write_boxes(self,
|
|
df: pd.DataFrame,
|
|
ts: VideoTimestamp,
|
|
frame_index: int) -> None:
|
|
if self._is_empty_df(df):
|
|
return
|
|
cur = self._conn.cursor()
|
|
res = cur.execute('''
|
|
INSERT INTO video_frame(timestamp_, index_, video_id, nbr_detected_object )
|
|
VALUES (?, ?, ?, 1)
|
|
ON CONFLICT(video_id, index_)
|
|
DO UPDATE SET nbr_detected_object = nbr_detected_object + 1
|
|
RETURNING id
|
|
''', (ts, frame_index, self._video_id))
|
|
frame_id = res.fetchone()[0]
|
|
cols = ['xmin', 'ymin', 'xmax', 'ymax', 'score', 'category', 'object_id']
|
|
cur.executemany('''
|
|
INSERT INTO tracked_detected_object(xmin,
|
|
ymin,
|
|
xmax,
|
|
ymax,
|
|
score,
|
|
category,
|
|
object_id,
|
|
video_frame_id)
|
|
VALUES (?, ?, ?, ?, ?, ?, ?, ?)
|
|
''', (list(x[cols]) + [frame_id] for x in df.iloc))
|
|
self._conn.commit()
|