db.py

import sqlite3
from datetime import datetime

DB_NAME="prediction.db"

def init_db():
    conn=sqlite3.connect(DB_NAME)
    cursor=conn.cursor()

    cursor.execute("""
        CREATE TABLE IF NOT EXISTS predictions (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            filename TEXT,
            top1_class TEXT,
            confidence REAL,
            latency_ms REAL,
            timestamp TEXT
        )
    """)
    conn.commit()
    conn.close()

def save_prediction(filename, top1_class, confidence, latency_ms):
    conn = sqlite3.connect(DB_NAME)
    cursor = conn.cursor()
    cursor.execute("""
        INSERT INTO predictions (
            filename,
            top1_class,
            confidence,
            latency_ms,
            timestamp
        )
        VALUES (?, ?, ?, ?, ?)
    """, (
        filename,
        top1_class,
        confidence,
        latency_ms,
        datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    ))
    conn.commit()
    conn.close()