This commit is contained in:
Zeecka
2025-12-27 00:08:28 +01:00
parent 1b62d80d47
commit 43938c410c
3 changed files with 43 additions and 42 deletions
+6 -4
View File
@@ -1,7 +1,7 @@
"""Image size bruteforcer for common PNG sizes."""
from pathlib import Path
from typing import List
from typing import Any, List
from ..app import app
from ..models import IHDR
@@ -73,8 +73,7 @@ def search_height_crc(png: PNG) -> List[PNG]:
results = []
for height in range(100, 3501):
candidate = PNG(
width=png.width,
height=height,
size=(png.width, height),
bit_depth=png.bit_depth,
color_type=png.color_type,
interlace=png.interlace,
@@ -125,7 +124,10 @@ def analyze_image_resize(input_img: Path, output_dir: Path) -> None:
try:
with input_img.open("rb") as image:
img_bytes = bytearray(image.read())
try:
png = PNG.from_bytearray(img_bytes)
except ValueError:
png = None
if png is None:
logs.append(
@@ -168,7 +170,7 @@ def analyze_image_resize(input_img: Path, output_dir: Path) -> None:
saved_img_urls.append("/image/" + str(Path(output_dir.name) / img_name))
# 4. Final Data Update
output_data = {
output_data: dict[str, Any] = {
"image_resize": {
"status": "ok",
"output": logs,
+16 -19
View File
@@ -9,7 +9,7 @@ import threading
import zlib
from functools import cached_property
from pathlib import Path
from typing import Any, Self
from typing import Any, Iterable, Self
_thread_lock = threading.Lock()
@@ -52,7 +52,7 @@ def update_data(
fcntl.flock(lock, fcntl.LOCK_UN)
def get_resolutions():
def get_resolutions() -> list[tuple[int, int]]:
"""
Generate a sorted list of common display and image resolutions.
@@ -104,7 +104,7 @@ def get_resolutions():
(14, 85), # Legal
(17, 11), # Tabloid
]
resolutions = set()
resolutions: set[tuple[int, int]] = set()
for w in base_widths:
for ar_w, ar_h in aspect_ratios:
h = math.floor(round(w * ar_h / ar_w))
@@ -114,7 +114,7 @@ def get_resolutions():
return sorted(resolutions)
def get_valid_depth_color_pairs():
def get_valid_depth_color_pairs() -> Iterable[tuple[int, int]]:
"""
Generate valid combinations of color and bit depth values.
@@ -154,20 +154,19 @@ class PNG:
crc: int
interlace: int = 0
def __init__(
def __init__( # pylint: disable=R0913
self,
width: int,
height: int,
size: tuple[int, int],
bit_depth: int,
color_type: int,
*,
crc: int | None = None,
interlace: int = 0,
):
"""Initialize a PNG IHDR (Image Header) chunk.
Args:
width (int): The width of the image in pixels.
height (int): The height of the image in pixels.
size (tuple[int, int]): The width and height of the image in pixels.
bit_depth (int): The number of bits per sample or per palette index (1, 2, 4, 8, or 16).
color_type (int): The color type of the image (0=Grayscale, 2=RGB, 3=Indexed,
4=Grayscale+Alpha, 6=RGBA).
@@ -177,16 +176,16 @@ class PNG:
interlace (int, optional): The interlace method (0=no interlace, 1=Adam7 interlace).
Defaults to 0.
"""
self.width = width
self.height = height
self.width = size[0]
self.height = size[1]
self.bit_depth = bit_depth
self.color_type = color_type
self.crc = (
crc
or zlib.crc32(
b"IHDR"
+ width.to_bytes(4, "big")
+ height.to_bytes(4, "big")
+ self.width.to_bytes(4, "big")
+ self.height.to_bytes(4, "big")
+ bytes((bit_depth, color_type, 0, 0, interlace))
)
& 0xFFFFFFFF
@@ -214,8 +213,7 @@ class PNG:
height = (packed >> 7) & 0xFFFF
width = (packed >> 23) & 0xFFFF
return cls(
width=width,
height=height,
size=(width, height),
bit_depth=depth,
color_type=color,
crc=crc,
@@ -267,14 +265,14 @@ class PNG:
"""
png_signature = b"\x89PNG\r\n\x1a\n"
if data[:8] != png_signature:
if bytes(data[:8]) != png_signature:
raise ValueError("Invalid PNG signature")
offset = 8 # Read IHDR chunk
length = struct.unpack(">I", data[offset : offset + 4])[0]
offset += 4
chunk_type = data[offset : offset + 4]
offset += 4
if chunk_type != b"IHDR":
if bytes(chunk_type) != b"IHDR":
raise ValueError("IHDR chunk not found at expected position")
ihdr_data = data[offset : offset + length]
@@ -288,8 +286,7 @@ class PNG:
)
return cls(
width=width,
height=height,
size=(width, height),
bit_depth=depth,
color_type=color,
crc=crc,
+20 -18
View File
@@ -15,21 +15,21 @@ db: SQLAlchemy = SQLAlchemy()
class Image(db.Model): # type: ignore
"""Model representing an image file in the database."""
hash: Column[String] = Column(
hash = Column(
String(64), primary_key=True, unique=True, nullable=False
)
file: Column[String] = Column(String(128), unique=True, nullable=False)
size: Column[Integer] = Column(Integer, nullable=False)
first_submission_date: Column[DateTime] = Column(
file = Column(String(128), unique=True, nullable=False)
size = Column(Integer, nullable=False)
first_submission_date = Column(
DateTime, nullable=False, default=lambda: datetime.now(timezone.utc)
)
last_submission_date: Column[DateTime] = Column(
last_submission_date = Column(
DateTime,
nullable=False,
default=lambda: datetime.now(timezone.utc),
onupdate=lambda: datetime.now(timezone.utc),
)
upload_count: Column[Integer] = Column(Integer)
upload_count = Column(Integer)
submissions = db.relationship("Submission", backref="image", lazy=True)
@@ -39,14 +39,14 @@ class Submission(db.Model): # type: ignore
option.
"""
hash: Column[String] = Column(
hash = Column(
String(128), primary_key=True, unique=True, nullable=False
)
filename: Column[String] = Column(String(128), nullable=False)
password: Column[String] = Column(String(128))
deep_analysis: Column[Boolean] = Column(Boolean, default=False)
status: Column[String] = Column(String(20), default="pending")
date: Column[Float] = Column(
filename = Column(String(128), nullable=False)
password = Column(String(128))
deep_analysis = Column(Boolean, default=False)
status = Column(String(20), default="pending")
date = Column(
Float, nullable=False, default=lambda: datetime.now(timezone.utc)
)
@@ -57,9 +57,9 @@ class Submission(db.Model): # type: ignore
class IHDR(db.Model): # type: ignore
"""IHDR CRC lookup table"""
iid: Column[Integer] = Column(Integer, primary_key=True, autoincrement=True)
crc: Column[BigInteger] = Column(BigInteger, nullable=False)
packed: Column[BigInteger] = Column(BigInteger, nullable=False)
iid = Column(Integer, primary_key=True, autoincrement=True)
crc = Column(BigInteger, nullable=False)
packed = Column(BigInteger, nullable=False)
def create_crc_db() -> None:
@@ -87,8 +87,8 @@ def create_crc_db() -> None:
count = 0
with db.session.no_autoflush: # pylint: disable=no-member
for (w, h), (bd, ct), inter in combinations:
p = PNG(width=w, height=h, bit_depth=bd, color_type=ct, interlace=inter)
for size, (bd, ct), inter in combinations:
p = PNG(size=size, bit_depth=bd, color_type=ct, interlace=inter)
exists = db.session.execute( # pylint: disable=no-member
db.select(IHDR.iid).where(IHDR.packed == p.packed)
@@ -97,7 +97,9 @@ def create_crc_db() -> None:
if exists: # Skip if already present
continue
db.session.add(IHDR(packed=p.packed, crc=p.crc)) # pylint: disable=no-member
db.session.add( # pylint: disable=no-member
IHDR(packed=p.packed, crc=p.crc)
)
count += 1
if count % 5000 == 0: