Spaces:
Running
Running
| """ | |
| 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 |