scripts/quick_analysis.py
| 1 | import os |
| 2 | from PIL import Image |
| 3 | from collections import Counter |
| 4 | |
| 5 | def analyze_images(directory): |
| 6 | analysis_results = {} |
| 7 | |
| 8 | for root, dirs, files in os.walk(directory): |
| 9 | if files: |
| 10 | model_folder_name = os.path.basename(root) |
| 11 | if model_folder_name not in analysis_results: |
| 12 | analysis_results[model_folder_name] = { |
| 13 | 'image_count': 0, |
| 14 | 'total_size': 0, |
| 15 | 'resolutions': Counter() |
| 16 | } |
| 17 | |
| 18 | for file in files: |
| 19 | file_path = os.path.join(root, file) |
| 20 | |
| 21 | # Count the image |
| 22 | analysis_results[model_folder_name]['image_count'] += 1 |
| 23 | |
| 24 | # Calculate the size of the image |
| 25 | try: |
| 26 | with Image.open(file_path) as img: |
| 27 | # Get the size of the image in bytes |
| 28 | file_size = os.path.getsize(file_path) |
| 29 | analysis_results[model_folder_name]['total_size'] += file_size |
| 30 | |
| 31 | # Get image dimensions |
| 32 | width, height = img.size |
| 33 | analysis_results[model_folder_name]['resolutions'][(width, height)] += 1 |
| 34 | except Exception as e: |
| 35 | print(f"Error reading file {file_path}: {e}") |
| 36 | |
| 37 | return analysis_results |
| 38 | |
| 39 | def print_and_log_analysis_results(analysis_results, dataset_name, log_file): |
| 40 | # Determine the maximum length of model names |
| 41 | max_model_length = max(len(model) for model in analysis_results.keys()) |
| 42 | model_column_width = max(max_model_length, 20) # Ensure at least 20 characters |
| 43 | |
| 44 | # Define column widths |
| 45 | image_count_width = 12 |
| 46 | total_size_width = 14 |
| 47 | resolution_width = 25 |
| 48 | |
| 49 | # Create header |
| 50 | header = f"{'Model':<{model_column_width}} | {'Image Count':>{image_count_width}} | {'Total Size (MB)':>{total_size_width}} | {'Most Common Resolution':<{resolution_width}}" |
| 51 | separator = "-" * (model_column_width + image_count_width + total_size_width + resolution_width + 7) # 7 for separators |
| 52 | |
| 53 | result_lines = [] |
| 54 | result_lines.append(f"Analysis for {dataset_name}:\n") |
| 55 | result_lines.append(header + "\n") |
| 56 | result_lines.append(separator + "\n") |
| 57 | |
| 58 | for model, data in analysis_results.items(): |
| 59 | total_size_mb = data['total_size'] / (1024 * 1024) |
| 60 | most_common_resolution = data['resolutions'].most_common(1) |
| 61 | |
| 62 | if most_common_resolution: |
| 63 | common_res = f"{most_common_resolution[0][0][0]}x{most_common_resolution[0][0][1]} ({most_common_resolution[0][1]} images)" |
| 64 | else: |
| 65 | common_res = "None" |
| 66 | |
| 67 | result_lines.append(f"{model:<{model_column_width}} | {data['image_count']:>{image_count_width}} | {total_size_mb:>{total_size_width}.2f} | {common_res:<{resolution_width}}\n") |
| 68 | |
| 69 | result_lines.append("\n") |
| 70 | |
| 71 | # Print to console |
| 72 | for line in result_lines: |
| 73 | print(line, end='') |
| 74 | |
| 75 | # Write to log file |
| 76 | with open(log_file, 'a') as f: |
| 77 | f.writelines(result_lines) |
| 78 | |
| 79 | def main(): |
| 80 | # Define directories |
| 81 | generated_dir = 'resampledEvalSet' |
| 82 | real_dir = 'real' |
| 83 | log_file = 'analysis_results.txt' |
| 84 | |
| 85 | # Clear the log file (optional, comment out if you want to append) |
| 86 | with open(log_file, 'w') as f: |
| 87 | pass |
| 88 | |
| 89 | # Analyze generated images |
| 90 | generated_analysis_results = analyze_images(generated_dir) |
| 91 | print_and_log_analysis_results(generated_analysis_results, "Generated Images", log_file) |
| 92 | |
| 93 | # Analyze real images |
| 94 | real_analysis_results = analyze_images(real_dir) |
| 95 | print_and_log_analysis_results(real_analysis_results, "Real Images", log_file) |
| 96 | |
| 97 | if __name__ == "__main__": |
| 98 | main() |
| 99 | |