121 lines
2.8 KiB
Python
121 lines
2.8 KiB
Python
import sqlite3
|
|
from pathlib import Path
|
|
|
|
|
|
class sql_database:
|
|
def __init__ (self, db_name="database.db"):
|
|
self.db_name = db_name
|
|
|
|
def connect(self):
|
|
"""Create a database connection."""
|
|
return sqlite3.connect(self.db_name)
|
|
|
|
def create_table(self):
|
|
"""Create table and structure."""
|
|
conn = self.connect()
|
|
|
|
conn.execute("""
|
|
CREATE TABLE IF NOT EXISTS image (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
file_name TEXT NOT NULL,
|
|
tag TEXT NOT NULL
|
|
)
|
|
""")
|
|
|
|
conn.commit()
|
|
conn.close()
|
|
|
|
def add_entry(self, entry_file_name, entry_tag):
|
|
"""Add new entry"""
|
|
conn = self.connect()
|
|
|
|
try:
|
|
conn.execute("""
|
|
INSERT INTO image (file_name, tag)
|
|
VALUES (?,?)
|
|
""",(entry_file_name, entry_tag))
|
|
|
|
conn.commit()
|
|
|
|
except sqlite3.IntegrityError as e:
|
|
print(f"Database error: {e}")
|
|
|
|
finally: conn.close()
|
|
|
|
def get_all_entry(self):
|
|
"""Get all entry."""
|
|
conn = self.connect()
|
|
|
|
cursor = conn.execute("""
|
|
SELECT id, file_name, tag
|
|
FROM image
|
|
""")
|
|
|
|
all_entry = cursor.fetchall()
|
|
|
|
conn.close()
|
|
|
|
return all_entry
|
|
|
|
def get_entry(self, entry_id):
|
|
"""Get an existing entry"""
|
|
conn = self.connect()
|
|
|
|
cursor = conn.execute("""
|
|
SELECT id, file_name,tag
|
|
FROM image
|
|
WHERE id = ?
|
|
""", (entry_id,))
|
|
|
|
entry = cursor.fetchone()
|
|
|
|
conn.close()
|
|
|
|
return entry
|
|
|
|
def update_entry(self, entry_id, entry_file_name, entry_tag):
|
|
"""Update an existing entry"""
|
|
conn = self.connect()
|
|
|
|
conn.execute("""
|
|
UPDATE image
|
|
SET file_name = ?, tag = ?
|
|
WHERE id = ?
|
|
""", (entry_file_name, entry_tag, entry_id))
|
|
|
|
conn.commit()
|
|
conn.close()
|
|
|
|
def delete_entry(self, entry_id):
|
|
"""Delete an entry"""
|
|
conn = self.connect()
|
|
conn.execute("""
|
|
DELETE FROM image
|
|
WHERE id = ?
|
|
""", (entry_id,))
|
|
|
|
conn.commit()
|
|
conn.close()
|
|
def test(self):
|
|
print()
|
|
|
|
def scan_folder(db:sql_database, dataset_path):
|
|
for folder in dataset_path.iterdir():
|
|
|
|
if not folder.is_dir():
|
|
continue
|
|
|
|
tag = folder.name
|
|
for file_path in folder.iterdir():
|
|
db.add_entry(entry_file_name=str(file_path),
|
|
entry_tag=tag)
|
|
if __name__ == "__main__":
|
|
grape_db = sql_database("grape.db")
|
|
|
|
grape_db.create_table()
|
|
|
|
dataset_path = Path("database/data")
|
|
|
|
scan_folder(grape_db, dataset_path)
|
|
|
|
|