This commit is contained in:
Philip Zimmermann
2025-12-13 16:37:46 +01:00
parent 9cacbf7d13
commit 705e19701f
2 changed files with 39 additions and 23 deletions
+10 -7
View File
@@ -3,6 +3,7 @@
import zlib import zlib
import sqlite3 import sqlite3
from pathlib import Path from pathlib import Path
from typing import List
from .utils import update_data from .utils import update_data
# Full list of expected sizes # Full list of expected sizes
@@ -182,7 +183,7 @@ def calc_checksum(
return bytearray((zlib.crc32(new_header) & 0xFFFFFFFF).to_bytes(4, byteorder="big")) return bytearray((zlib.crc32(new_header) & 0xFFFFFFFF).to_bytes(4, byteorder="big"))
def lookup_crc(crc_bytes: bytes, logs: list) -> list: def lookup_crc(crc_bytes: bytes, logs: List[str]) -> List[tuple[int, int]]:
""" """
Queries the SQLite DB for the CRC and returns a list of (width, height) tuples. Queries the SQLite DB for the CRC and returns a list of (width, height) tuples.
""" """
@@ -196,7 +197,7 @@ def lookup_crc(crc_bytes: bytes, logs: list) -> list:
try: try:
target_crc = int.from_bytes(crc_bytes, byteorder="big") target_crc = int.from_bytes(crc_bytes, byteorder="big")
except ValueError: except ValueError:
logs.append(f"Invalid CRC bytes: {crc_bytes}") logs.append(f"Invalid CRC bytes: {crc_bytes!r}")
return [] return []
conn = sqlite3.connect(db_path) conn = sqlite3.connect(db_path)
@@ -210,9 +211,9 @@ def lookup_crc(crc_bytes: bytes, logs: list) -> list:
if results: if results:
logs.append(f"Database: Found {len(results)} match(es) for CRC {target_crc}") logs.append(f"Database: Found {len(results)} match(es) for CRC {target_crc}")
return results return results
else:
logs.append(f"Database: No match found for CRC {target_crc}") logs.append(f"Database: No match found for CRC {target_crc}")
return [] return []
def analyze_image_resize(input_img: Path, output_dir: Path) -> None: def analyze_image_resize(input_img: Path, output_dir: Path) -> None:
@@ -269,7 +270,9 @@ def analyze_image_resize(input_img: Path, output_dir: Path) -> None:
# Check if this size matches the file's CRC # Check if this size matches the file's CRC
if target_crc == calc_checksum(header_chunk, width_bytes, height_bytes): if target_crc == calc_checksum(header_chunk, width_bytes, height_bytes):
match_found = True match_found = True
logs.append(f"Success (Common List): Match found! Dimensions: {width}x{height}") logs.append(
f"Success (Common List): Match found! Dimensions: {width}x{height}"
)
candidates.append((width, height)) candidates.append((width, height))
break # Stop after first match break # Stop after first match
@@ -314,4 +317,4 @@ def analyze_image_resize(input_img: Path, output_dir: Path) -> None:
update_data(output_dir, output_data) update_data(output_dir, output_data)
except Exception as e: except Exception as e:
update_data(output_dir, {"image_resize": {"status": "error", "error": str(e)}}) update_data(output_dir, {"image_resize": {"status": "error", "error": str(e)}})
+29 -16
View File
@@ -14,6 +14,14 @@ _thread_lock = threading.Lock()
MAX_PENDING_TIME = int(os.getenv("MAX_PENDING_TIME", "600")) # 10 minutes by default MAX_PENDING_TIME = int(os.getenv("MAX_PENDING_TIME", "600")) # 10 minutes by default
# PATH CONFIGURATION
# Prioritize /data for Docker persistence, fallback to local directory for bare metal
DB_NAME = "ihdr_crcs.db"
if os.path.exists("/data"):
DB_PATH = os.path.join("/data", DB_NAME)
else:
DB_PATH = os.path.join(os.path.dirname(os.path.abspath(__file__)), DB_NAME)
def update_data( def update_data(
output_dir: Path, new_data: dict[Any, Any], json_filename: str = "results.json" output_dir: Path, new_data: dict[Any, Any], json_filename: str = "results.json"
@@ -51,7 +59,11 @@ def update_data(
fcntl.flock(lock, fcntl.LOCK_UN) fcntl.flock(lock, fcntl.LOCK_UN)
def create_crc_db(): def create_crc_db() -> None:
"""Creates and populates CRC lookup database with CRC, Width and Heigth tuples"""
if os.path.exists(DB_PATH):
print(f"Database already exists at {DB_PATH}. Skipping generation.")
return
db_filename = "ihdr_crcs.db" db_filename = "ihdr_crcs.db"
# Connect to database # Connect to database
@@ -64,51 +76,53 @@ def create_crc_db():
# Create table # Create table
# We store the CRC as an integer for faster indexing/storage # We store the CRC as an integer for faster indexing/storage
cursor.execute(""" cursor.execute(
"""
CREATE TABLE IF NOT EXISTS ihdr ( CREATE TABLE IF NOT EXISTS ihdr (
crc INTEGER, crc INTEGER,
width INTEGER, width INTEGER,
height INTEGER height INTEGER
) )
""") """
)
# Drop index if exists (rebuilding it at the end is faster) # Drop index if exists (rebuilding it at the end is faster)
cursor.execute("DROP INDEX IF EXISTS idx_crc") cursor.execute("DROP INDEX IF EXISTS idx_crc")
max_width = 2000 max_width = 2000
max_height = 2000 max_height = 2000
W = range(1, max_width + 1) w = range(1, max_width + 1)
H = range(1, max_height + 1) h = range(1, max_height + 1)
# Standard PNG IHDR parameters # Standard PNG IHDR parameters
bit_depth = [1, 2, 4, 8, 16] bit_depth = [1, 2, 4, 8, 16]
color_type = [0, 2, 3, 4, 6] color_type = [0, 2, 3, 4, 6]
compression_method = [0] # This is always 0 according to the PNG spec compression_method = [0] # This is always 0 according to the PNG spec
filter_method = [0] # This is always 0 according to the PNG spec filter_method = [0] # This is always 0 according to the PNG spec
interlace_method = [0, 1] interlace_method = [0, 1]
# Create the generator # Create the generator
all_combinations = itertools.product( all_combinations = itertools.product(
W, H, bit_depth, color_type, compression_method, filter_method, interlace_method w, h, bit_depth, color_type, compression_method, filter_method, interlace_method
) )
batch_size = 100000 batch_size = 100000
batch = [] batch = []
count = 0 count = 0
for w, h, bd, ct, comp, filt, inter in all_combinations: for w_val, h_val, bd, ct, comp, filt, inter in all_combinations:
# Construct IHDR chunk data # Construct IHDR chunk data
ihdr_data = ( ihdr_data = (
w.to_bytes(4, "big") + w_val.to_bytes(4, "big")
h.to_bytes(4, "big") + + h_val.to_bytes(4, "big")
bytes([bd, ct, comp, filt, inter]) + bytes([bd, ct, comp, filt, inter])
) )
# Calculate CRC # Calculate CRC
crc = zlib.crc32(b"IHDR" + ihdr_data) & 0xffffffff crc = zlib.crc32(b"IHDR" + ihdr_data) & 0xFFFFFFFF
# Add to batch # Add to batch
batch.append((crc, w, h)) batch.append((crc, w_val, h_val))
if len(batch) >= batch_size: if len(batch) >= batch_size:
cursor.executemany("INSERT INTO ihdr VALUES (?,?,?)", batch) cursor.executemany("INSERT INTO ihdr VALUES (?,?,?)", batch)
@@ -129,5 +143,4 @@ def create_crc_db():
conn.commit() conn.commit()
conn.close() conn.close()
print(f"Database ready.") print("Database ready.")
return