Spaces:
Running
Running
File size: 4,265 Bytes
31e20c4 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 |
"""
File utilities for Image Tagger application.
"""
import os
import time
def save_tags_to_file(image_path, all_tags, original_filename=None, custom_dir=None, overwrite=False):
"""
Save tags to a text file in a dedicated 'saved_tags' folder or custom directory.
Args:
image_path: Path to the original image
all_tags: List of all tags to save
original_filename: Original filename if uploaded through Streamlit
custom_dir: Custom directory to save tags to (if None, uses 'saved_tags' folder)
Returns:
Path to the saved file
"""
# Determine the save directory
if custom_dir and os.path.isdir(custom_dir):
save_dir = custom_dir
else:
# Create a dedicated folder for saved tags in the app's root directory
app_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
save_dir = os.path.join(app_dir, "saved_tags")
# Ensure the directory exists
os.makedirs(save_dir, exist_ok=True)
# Determine the filename
if original_filename:
# For uploaded files, use original filename
base_name = os.path.splitext(original_filename)[0]
else:
# For non-uploaded files, use the image path
base_name = os.path.splitext(os.path.basename(image_path))[0]
# Create the output path
output_path = os.path.join(save_dir, f"{base_name}.txt")
# If overwrite is False and file exists, add a timestamp to avoid overwriting
if not overwrite and os.path.exists(output_path):
timestamp = time.strftime("%Y%m%d-%H%M%S")
output_path = os.path.join(save_dir, f"{base_name}_{timestamp}.txt")
# Write the tags to file
with open(output_path, 'w', encoding='utf-8') as f:
if all_tags:
# Add comma after each tag including the last one
tag_text = ", ".join(all_tags) + ","
f.write(tag_text)
return output_path
def get_default_save_locations():
"""
Get default save locations for tag files.
Returns:
List of default save locations
"""
# App directory
app_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
save_dir = os.path.join(app_dir, "saved_tags")
# Common user directories
desktop_dir = os.path.expanduser("~/Desktop")
download_dir = os.path.expanduser("~/Downloads")
documents_dir = os.path.expanduser("~/Documents")
# List of default save locations
save_locations = [
save_dir,
desktop_dir,
download_dir,
documents_dir,
]
# Ensure directories exist
for folder in save_locations:
os.makedirs(folder, exist_ok=True)
return save_locations
def apply_category_limits(result, category_limits):
"""
Apply category limits to a result dictionary.
Args:
result: Result dictionary containing tags and all_tags
category_limits: Dictionary mapping categories to their tag limits
(0 = exclude category, -1 = no limit/include all)
Returns:
Updated result dictionary with limits applied
"""
if not category_limits or not result['success']:
return result
# Get the filtered tags
filtered_tags = result['tags']
# Apply limits to each category
for category, cat_tags in list(filtered_tags.items()):
# Get limit for this category, default to -1 (no limit)
limit = category_limits.get(category, -1)
if limit == 0:
# Exclude this category entirely
del filtered_tags[category]
elif limit > 0 and len(cat_tags) > limit:
# Limit to top N tags for this category
filtered_tags[category] = cat_tags[:limit]
# Regenerate all_tags list after applying limits
all_tags = []
for category, cat_tags in filtered_tags.items():
for tag, _ in cat_tags:
all_tags.append(tag)
# Update the result with limited tags
result['tags'] = filtered_tags
result['all_tags'] = all_tags
return result |