75 lines
2.9 KiB
Python
75 lines
2.9 KiB
Python
|
"""
|
||
|
|
Script to iterate over all images in the img folder and insert their metadata into the avdb_image table.
|
||
|
|
"""
|
||
|
|
import os
|
||
|
|
from PIL import Image
|
||
|
|
import sqlalchemy
|
||
|
|
from sqlalchemy import text, exc
|
||
|
|
|
||
|
|
# Database URL from alembic.ini
|
||
|
|
DATABASE_URL = "postgresql+psycopg2://postgres:pgRocks@localhost:5432/testdb"
|
||
|
|
|
||
|
|
def get_image_metadata(image_path):
|
||
|
|
"""Extract metadata (width, height) from an image."""
|
||
|
|
with Image.open(image_path) as img:
|
||
|
|
width, height = img.size
|
||
|
|
return width, height
|
||
|
|
|
||
|
|
def insert_images_into_db():
|
||
|
|
"""Iterate over all images in the img folder and insert their metadata into the avdb_image table."""
|
||
|
|
engine = sqlalchemy.create_engine(DATABASE_URL)
|
||
|
|
|
||
|
|
# Iterate over all subdirectories in the img folder
|
||
|
|
img_folder = "img"
|
||
|
|
inserted_count = 0
|
||
|
|
skipped_count = 0
|
||
|
|
test_shim = 'TESTIMG.jpeg'
|
||
|
|
duplicate_count = 0
|
||
|
|
|
||
|
|
for root, dirs, files in os.walk(img_folder):
|
||
|
|
for file in files:
|
||
|
|
# test nmode
|
||
|
|
if file.endswith(test_shim):
|
||
|
|
# if file.lower().endswith(('.jpg', '.jpeg')):
|
||
|
|
image_path = os.path.join(root, file)
|
||
|
|
img_name = file
|
||
|
|
|
||
|
|
# Extract width and height
|
||
|
|
try:
|
||
|
|
width, height = get_image_metadata(image_path)
|
||
|
|
except Exception as e:
|
||
|
|
print(f"Failed to extract metadata for {image_path}: {e}")
|
||
|
|
skipped_count += 1
|
||
|
|
continue
|
||
|
|
|
||
|
|
# Insert into the database
|
||
|
|
try:
|
||
|
|
with engine.connect() as connection:
|
||
|
|
query = text(
|
||
|
|
"""
|
||
|
|
INSERT INTO avdb_image (img_path, img_name, width, height)
|
||
|
|
VALUES (:img_path, :img_name, :width, :height)
|
||
|
|
"""
|
||
|
|
)
|
||
|
|
connection.execute(
|
||
|
|
query,
|
||
|
|
{"img_path": image_path, "img_name": img_name, "width": width, "height": height}
|
||
|
|
)
|
||
|
|
connection.commit()
|
||
|
|
inserted_count += 1
|
||
|
|
print(f"Inserted: {img_name} (Width: {width}, Height: {height})")
|
||
|
|
except exc.IntegrityError:
|
||
|
|
# Handle duplicate key violations
|
||
|
|
print(f"Duplicate entry skipped for {image_path}")
|
||
|
|
duplicate_count += 1
|
||
|
|
except Exception as e:
|
||
|
|
print(f"Failed to insert {image_path} into the database: {e}")
|
||
|
|
skipped_count += 1
|
||
|
|
|
||
|
|
print(f"\nTotal images inserted: {inserted_count}")
|
||
|
|
print(f"Total duplicates skipped: {duplicate_count}")
|
||
|
|
print(f"Total images skipped due to errors: {skipped_count}")
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
insert_images_into_db()
|