wehub-resource-sync
73 lines
2.9 KiB
Python
73 lines
2.9 KiB
Python
# @Author: Pevernow (wzy3450354617@gmail.com)
|
|
# @Date: 2025/1/5
|
|
# @License: (Follow the main project)
|
|
import json
|
|
import os
|
|
import tarfile
|
|
|
|
from PIL import Image, PngImagePlugin
|
|
|
|
PngImagePlugin.MAX_TEXT_CHUNK = 100 * 1024 * 1024 # Increase maximum size for text chunks
|
|
|
|
|
|
def process_data(input_dir, output_tar_name="output.tar"):
|
|
"""
|
|
Processes a directory containing PNG files, generates corresponding JSON files,
|
|
and packages all files into a TAR file. It also counts the number of processed PNG images,
|
|
and saves the height and width of each PNG file to the JSON.
|
|
|
|
Args:
|
|
input_dir (str): The input directory containing PNG files.
|
|
output_tar_name (str): The name of the output TAR file (default is "output.tar").
|
|
"""
|
|
png_count = 0
|
|
json_files_created = []
|
|
|
|
for filename in os.listdir(input_dir):
|
|
if filename.lower().endswith(".png"):
|
|
png_count += 1
|
|
base_name = filename[:-4] # Remove the ".png" extension
|
|
txt_filename = os.path.join(input_dir, base_name + ".txt")
|
|
json_filename = base_name + ".json"
|
|
json_filepath = os.path.join(input_dir, json_filename)
|
|
png_filepath = os.path.join(input_dir, filename)
|
|
|
|
if os.path.exists(txt_filename):
|
|
try:
|
|
# Get the dimensions of the PNG image
|
|
with Image.open(png_filepath) as img:
|
|
width, height = img.size
|
|
|
|
with open(txt_filename, encoding="utf-8") as f:
|
|
caption_content = f.read().strip()
|
|
|
|
data = {"file_name": filename, "prompt": caption_content, "width": width, "height": height}
|
|
|
|
with open(json_filepath, "w", encoding="utf-8") as outfile:
|
|
json.dump(data, outfile, indent=4, ensure_ascii=False)
|
|
|
|
print(f"Generated: {json_filename}")
|
|
json_files_created.append(json_filepath)
|
|
|
|
except Exception as e:
|
|
print(f"Error processing file {filename}: {e}")
|
|
else:
|
|
print(f"Warning: No corresponding TXT file found for {filename}.")
|
|
|
|
# Create a TAR file and include all files
|
|
with tarfile.open(output_tar_name, "w") as tar:
|
|
for item in os.listdir(input_dir):
|
|
item_path = os.path.join(input_dir, item)
|
|
tar.add(item_path, arcname=item) # arcname maintains the relative path of the file in the tar
|
|
|
|
print(f"\nAll files have been packaged into: {output_tar_name}")
|
|
print(f"Number of PNG images processed: {png_count}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
input_directory = input("Please enter the directory path containing PNG and TXT files: ")
|
|
output_tar_filename = (
|
|
input("Please enter the name of the output TAR file (default is output.tar): ") or "output.tar"
|
|
)
|
|
process_data(input_directory, output_tar_filename)
|