Files
novafarma/scripts/utils/perfect_cleanup.py
2026-01-25 12:20:50 +01:00

101 lines
3.7 KiB
Python

import os
import numpy as np
from PIL import Image
# MAPPING: Uploaded File -> Target Name
file_map = {
"uploaded_image_0_1768981117686.png": "gumb_start.png",
"uploaded_image_1_1768981117686.png": "merilec_zdravja.png",
"uploaded_image_2_1768981117686.png": "okvir_zarjavel.png",
"uploaded_image_3_1768981117686.jpg": "amnezija_maska.png"
}
# DIRECTORIES
upload_dir = "/Users/davidkotnik/.gemini/antigravity/brain/998d0b10-1733-4e5b-85ed-249b986ba9b3"
dest_dir = "/Users/davidkotnik/repos/novafarma/assets/slike/NOVE_SLIKE/UI"
if not os.path.exists(dest_dir):
os.makedirs(dest_dir)
def clean_image(img_path, output_path, add_green_bg=False):
try:
if not os.path.exists(img_path):
print(f"Skipping missing: {img_path}")
return
img = Image.open(img_path).convert("RGBA")
data = np.array(img)
# Color Channels
r, g, b, a = data[:,:,0], data[:,:,1], data[:,:,2], data[:,:,3]
# --- CHECKERBOARD DETECTION LOGIC ---
# 1. Pure White/Near White
white_mask = (r > 230) & (g > 230) & (b > 230)
# 2. Pure Grey/Near Grey
# Standard checkerboard grey is often exactly equal R=G=B, around 204 or 192.
# We look for low saturation (R~=G~=B) and specific brightness range.
# Strict neutrality check (R, G, B within 10 of each other)
is_neutral = (np.abs(r.astype(int) - g.astype(int)) < 15) & \
(np.abs(r.astype(int) - b.astype(int)) < 15)
# Brightness range for typical checkerboards (Light Grey)
is_grey_range = (r > 160) & (r < 225)
grey_mask = is_neutral & is_grey_range
# 3. Green Screen (if any previous attempts left green)
green_mask = (g > 200) & (r < 100) & (b < 100)
# Combined Mask of "Pixels to Remove"
remove_mask = white_mask | grey_mask | green_mask
# --- APPLY CLEANING ---
# Set Alpha to 0 where mask matches
data[remove_mask, 3] = 0
# --- OPTIONAL: GREEN BACKGROUND ---
if add_green_bg:
# Create a solid green image
green_bg = np.zeros_like(data)
green_bg[:,:] = [0, 255, 0, 255] # Solid Green
# Where the original is NOT transparent, copy the original pixels
# We use the INVERSE of the remove_mask (pixels we kept)
# AND the original alpha channel (if original was significantly opaque)
# Simplified: Composite original over green
# But we just nuked the alpha of checkerboard, so we can standard composite
cleaned_img = Image.fromarray(data)
bg_img = Image.new("RGBA", cleaned_img.size, (0, 255, 0, 255))
bg_img.alpha_composite(cleaned_img)
# Save Green Version
bg_img.save(output_path)
print(f"Saved (GREEN BG): {output_path}")
else:
# Save Transparent Version
processed_img = Image.fromarray(data)
processed_img.save(output_path)
print(f"Saved (TRANSPARENT): {output_path}")
except Exception as e:
print(f"Error processing {img_path}: {e}")
# EXECUTE
for u_file, target_name in file_map.items():
src = os.path.join(upload_dir, u_file)
# 1. Save Transparent Version (Standard game use)
dst_trans = os.path.join(dest_dir, target_name)
clean_image(src, dst_trans, add_green_bg=False)
# 2. Save Green Version (As requested by user for manual keying?)
# dst_green = os.path.join(dest_dir, "green_" + target_name)
# clean_image(src, dst_green, add_green_bg=True)