Files
novafarma/scripts/test_ai_bg_removal.py

190 lines
5.4 KiB
Python

#!/usr/bin/env python3
"""
AI Background Removal using rembg
This will properly remove ANY background and create clean alpha transparency
"""
import os
import shutil
from pathlib import Path
from rembg import remove
from PIL import Image
# Test images
TEST_IMAGES = [
"assets/PHASE_PACKS/1_FAZA_1/tools/wood/watering_can.png",
"assets/PHASE_PACKS/1_FAZA_1/animals/horse.png",
"assets/PHASE_PACKS/1_FAZA_1/infrastructure/farm_elements/manure_pile.png",
"assets/PHASE_PACKS/1_FAZA_1/tools/iron/pickaxe.png",
"assets/PHASE_PACKS/1_FAZA_1/animals/sheep/walk.png",
]
OUTPUT_DIR = "test_transparency"
def ai_remove_bg(input_path, output_path):
"""Use AI (rembg) to remove background"""
img = Image.open(input_path)
# Remove background using AI
result = remove(img)
result.save(output_path, 'PNG')
print(f" ✅ AI Processed: {os.path.basename(output_path)}")
def main():
print("🤖 AI BACKGROUND REMOVAL (rembg)")
print("=" * 50)
print("\nUsing neural network to detect and remove backgrounds\n")
os.makedirs(OUTPUT_DIR, exist_ok=True)
processed = []
for img_path in TEST_IMAGES:
if not os.path.exists(img_path):
print(f" ❌ Not found: {img_path}")
continue
name = os.path.basename(img_path)
name_no_ext = os.path.splitext(name)[0]
# AI process
ai_dest = os.path.join(OUTPUT_DIR, f"{name_no_ext}_AI.png")
ai_remove_bg(img_path, ai_dest)
processed.append({
'name': name_no_ext,
'original': f"{name_no_ext}_ORIGINAL.png",
'ai': f"{name_no_ext}_AI.png"
})
# Generate comparison HTML
html = '''<!DOCTYPE html>
<html>
<head>
<meta charset="UTF-8">
<title>🤖 AI Background Removal Test</title>
<style>
body {
background: #1a1a2e;
color: white;
font-family: Arial, sans-serif;
padding: 30px;
}
h1 { color: #ff4444; text-align: center; }
.ref-note {
background: linear-gradient(135deg, #00ff00, #00aa00);
color: black;
padding: 20px;
border-radius: 10px;
text-align: center;
margin-bottom: 30px;
font-weight: bold;
}
.comparison {
display: flex;
gap: 30px;
margin: 30px 0;
padding: 20px;
background: rgba(0,0,0,0.3);
border-radius: 15px;
align-items: center;
}
.image-box {
flex: 1;
text-align: center;
}
.image-box h3 { color: #ff6666; margin-bottom: 15px; }
.image-box img {
max-width: 250px;
max-height: 250px;
border-radius: 10px;
}
.bg-checker {
background: linear-gradient(45deg, #333 25%, #555 25%, #555 50%, #333 50%, #333 75%, #555 75%);
background-size: 16px 16px;
padding: 15px;
border-radius: 10px;
}
.bg-green {
background: #00ff00;
padding: 15px;
border-radius: 10px;
}
.bg-game {
background: linear-gradient(to bottom, #3d6b1f, #234010);
padding: 15px;
border-radius: 10px;
}
.bg-red {
background: #cc3333;
padding: 15px;
border-radius: 10px;
}
.label { margin-top: 8px; font-size: 0.8em; color: #888; }
.success { color: #44ff44; }
</style>
</head>
<body>
<h1>🤖 AI Background Removal (rembg)</h1>
<div class="ref-note">
🎯 Using neural networks to detect subjects and remove backgrounds!
</div>
'''
for item in processed:
html += f'''
<div class="comparison">
<div class="image-box">
<h3>📷 ORIGINAL</h3>
<div class="bg-checker">
<img src="{item['original']}" alt="Original">
</div>
<div class="label">Original</div>
</div>
<div class="image-box">
<h3 class="success">🤖 AI (Checker)</h3>
<div class="bg-checker">
<img src="{item['ai']}" alt="AI">
</div>
<div class="label">Transparency check</div>
</div>
<div class="image-box">
<h3 class="success">🤖 AI (Green)</h3>
<div class="bg-green">
<img src="{item['ai']}" alt="On Green">
</div>
<div class="label">Green screen test</div>
</div>
<div class="image-box">
<h3 class="success">🤖 AI (Grass)</h3>
<div class="bg-game">
<img src="{item['ai']}" alt="In Game">
</div>
<div class="label">In-game preview</div>
</div>
<div class="image-box">
<h3 class="success">🤖 AI (Red)</h3>
<div class="bg-red">
<img src="{item['ai']}" alt="On Red">
</div>
<div class="label">Contrast test</div>
</div>
</div>
'''
html += '''
</body>
</html>
'''
html_path = os.path.join(OUTPUT_DIR, "comparison_ai.html")
with open(html_path, 'w') as f:
f.write(html)
print("\n" + "=" * 50)
print(f"✅ DONE! Open: {html_path}")
if __name__ == '__main__':
main()