Animal-Grid-Vectorizer / tests /test_background_removal.py
MakiAi's picture
Upload 84 files
ad8cacf verified
import os
from utils.background_remover import remove_background_simple, remove_background_advanced
# ใƒ†ใ‚นใƒˆ็”จใƒ‡ใ‚ฃใƒฌใ‚ฏใƒˆใƒชใฎไฝœๆˆ
test_output_dir = os.path.join('static', 'test_bg_removal')
os.makedirs(test_output_dir, exist_ok=True)
# ใƒ†ใ‚นใƒˆ็”ปๅƒใฎใƒ‘ใ‚น
test_img_dir = os.path.join('static', 'test_output', 'refined_animals')
if os.path.exists(test_img_dir):
print(f'ใƒ†ใ‚นใƒˆใƒ‡ใ‚ฃใƒฌใ‚ฏใƒˆใƒชใŒ่ฆ‹ใคใ‹ใ‚Šใพใ—ใŸ: {test_img_dir}')
# ใƒ‡ใ‚ฃใƒฌใ‚ฏใƒˆใƒชๅ†…ใฎ็”ปๅƒใ‚’ๅ–ๅพ—
image_files = [f for f in os.listdir(test_img_dir)
if f.endswith(('.png', '.jpg', '.jpeg'))]
if image_files:
test_img = os.path.join(test_img_dir, image_files[0])
print(f'ใƒ†ใ‚นใƒˆ็”ปๅƒใŒ่ฆ‹ใคใ‹ใ‚Šใพใ—ใŸ: {test_img}')
# ใ‚ทใƒณใƒ—ใƒซใช่ƒŒๆ™ฏ้™คๅŽปใ‚’ใƒ†ใ‚นใƒˆ
simple_output = os.path.join(test_output_dir, 'bg_removed_simple.png')
simple_result = remove_background_simple(test_img, simple_output)
print(f'ใ‚ทใƒณใƒ—ใƒซ่ƒŒๆ™ฏ้™คๅŽป็ตๆžœ: {simple_result}')
# ้ซ˜ๅบฆใช่ƒŒๆ™ฏ้™คๅŽปใ‚’ใƒ†ใ‚นใƒˆ
advanced_output = os.path.join(test_output_dir, 'bg_removed_advanced.png')
advanced_result = remove_background_advanced(test_img, advanced_output)
print(f'้ซ˜ๅบฆใช่ƒŒๆ™ฏ้™คๅŽป็ตๆžœ: {advanced_result}')
print('่ƒŒๆ™ฏ้™คๅŽปใƒ†ใ‚นใƒˆๅฎŒไบ†')
else:
print('ใƒ†ใ‚นใƒˆ็”ปๅƒใŒ่ฆ‹ใคใ‹ใ‚Šใพใ›ใ‚“ใงใ—ใŸ')
else:
print(f'ใƒ†ใ‚นใƒˆใƒ‡ใ‚ฃใƒฌใ‚ฏใƒˆใƒชใŒ่ฆ‹ใคใ‹ใ‚Šใพใ›ใ‚“ใงใ—ใŸ: {test_img_dir}')
# ใ‚ตใƒณใƒ—ใƒซ็”ปๅƒใƒ‡ใ‚ฃใƒฌใ‚ฏใƒˆใƒชใ‚’็ขบ่ช
sample_dir = 'static'
if os.path.exists(sample_dir):
print(f'ใ‚ตใƒณใƒ—ใƒซใƒ‡ใ‚ฃใƒฌใ‚ฏใƒˆใƒชใŒ่ฆ‹ใคใ‹ใ‚Šใพใ—ใŸ: {sample_dir}')
# ใ‚ตใƒณใƒ—ใƒซ็”ปๅƒใ‚’ๆŽขใ™
for root, dirs, files in os.walk(sample_dir):
image_files = [f for f in files if f.endswith(('.png', '.jpg', '.jpeg'))]
if image_files:
test_img = os.path.join(root, image_files[0])
print(f'ใ‚ตใƒณใƒ—ใƒซ็”ปๅƒใŒ่ฆ‹ใคใ‹ใ‚Šใพใ—ใŸ: {test_img}')
# ใ‚ทใƒณใƒ—ใƒซใช่ƒŒๆ™ฏ้™คๅŽปใ‚’ใƒ†ใ‚นใƒˆ
simple_output = os.path.join(test_output_dir, 'bg_removed_simple.png')
simple_result = remove_background_simple(test_img, simple_output)
print(f'ใ‚ทใƒณใƒ—ใƒซ่ƒŒๆ™ฏ้™คๅŽป็ตๆžœ: {simple_result}')
print('่ƒŒๆ™ฏ้™คๅŽปใƒ†ใ‚นใƒˆๅฎŒไบ†')
break
else:
print('ใ‚ตใƒณใƒ—ใƒซ็”ปๅƒใŒ่ฆ‹ใคใ‹ใ‚Šใพใ›ใ‚“ใงใ—ใŸ')