import os import shutil import tempfile import base64 import asyncio from io import BytesIO import cv2 import numpy as np import torch import onnxruntime as rt from PIL import Image import gradio as gr from transformers import pipeline from huggingface_hub import hf_hub_download # Import necessary function from aesthetic_predictor_v2_5 from aesthetic_predictor_v2_5 import convert_v2_5_from_siglip ##################################### # Model Definitions # ##################################### class MLP(torch.nn.Module): """A simple multi-layer perceptron for image feature regression.""" def __init__(self, input_size: int, batch_norm: bool = True): super().__init__() self.input_size = input_size self.layers = torch.nn.Sequential( torch.nn.Linear(self.input_size, 2048), torch.nn.ReLU(), torch.nn.BatchNorm1d(2048) if batch_norm else torch.nn.Identity(), torch.nn.Dropout(0.3), torch.nn.Linear(2048, 512), torch.nn.ReLU(), torch.nn.BatchNorm1d(512) if batch_norm else torch.nn.Identity(), torch.nn.Dropout(0.3), torch.nn.Linear(512, 256), torch.nn.ReLU(), torch.nn.BatchNorm1d(256) if batch_norm else torch.nn.Identity(), torch.nn.Dropout(0.2), torch.nn.Linear(256, 128), torch.nn.ReLU(), torch.nn.BatchNorm1d(128) if batch_norm else torch.nn.Identity(), torch.nn.Dropout(0.1), torch.nn.Linear(128, 32), torch.nn.ReLU(), torch.nn.Linear(32, 1) ) def forward(self, x: torch.Tensor) -> torch.Tensor: return self.layers(x) class WaifuScorer: """WaifuScorer model that uses CLIP for feature extraction and a custom MLP for scoring.""" def __init__(self, model_path: str = None, device: str = 'cuda', cache_dir: str = None, verbose: bool = False): self.verbose = verbose self.device = device self.dtype = torch.float32 self.available = False try: import clip # local import to avoid dependency issues # Set default model path if not provided if model_path is None: model_path = "Eugeoter/waifu-scorer-v3/model.pth" if self.verbose: print(f"Model path not provided. Using default: {model_path}") # Download model if not found locally if not os.path.isfile(model_path): username, repo_id, model_name = model_path.split("/")[-3:] model_path = hf_hub_download(f"{username}/{repo_id}", model_name, cache_dir=cache_dir) if self.verbose: print(f"Loading WaifuScorer model from: {model_path}") # Initialize MLP model self.mlp = MLP(input_size=768) # Load state dict if model_path.endswith(".safetensors"): from safetensors.torch import load_file state_dict = load_file(model_path) else: state_dict = torch.load(model_path, map_location=device) self.mlp.load_state_dict(state_dict) self.mlp.to(device) self.mlp.eval() # Load CLIP model for image preprocessing and feature extraction self.clip_model, self.preprocess = clip.load("ViT-L/14", device=device) self.available = True except Exception as e: print(f"Unable to initialize WaifuScorer: {e}") @torch.no_grad() def __call__(self, images): if not self.available: return [None] * (len(images) if isinstance(images, list) else 1) if isinstance(images, Image.Image): images = [images] n = len(images) # Ensure at least two images for CLIP model compatibility if n == 1: images = images * 2 image_tensors = [self.preprocess(img).unsqueeze(0) for img in images] image_batch = torch.cat(image_tensors).to(self.device) image_features = self.clip_model.encode_image(image_batch) # Normalize features norm = image_features.norm(2, dim=-1, keepdim=True) norm[norm == 0] = 1 im_emb = (image_features / norm).to(device=self.device, dtype=self.dtype) predictions = self.mlp(im_emb) scores = predictions.clamp(0, 10).cpu().numpy().reshape(-1).tolist() return scores[:n] ##################################### # Aesthetic Predictor Functions # ##################################### def load_aesthetic_predictor_v2_5(): """Load and return an instance of Aesthetic Predictor V2.5 with batch processing support.""" class AestheticPredictorV2_5_Impl: def __init__(self): print("Loading Aesthetic Predictor V2.5...") self.model, self.preprocessor = convert_v2_5_from_siglip( low_cpu_mem_usage=True, trust_remote_code=True, ) if torch.cuda.is_available(): self.model = self.model.to(torch.bfloat16).cuda() def inference(self, image): if isinstance(image, list): images_rgb = [img.convert("RGB") for img in image] pixel_values = self.preprocessor(images=images_rgb, return_tensors="pt").pixel_values if torch.cuda.is_available(): pixel_values = pixel_values.to(torch.bfloat16).cuda() with torch.inference_mode(): scores = self.model(pixel_values).logits.squeeze().float().cpu().numpy() if scores.ndim == 0: scores = np.array([scores]) return scores.tolist() else: pixel_values = self.preprocessor(images=image.convert("RGB"), return_tensors="pt").pixel_values if torch.cuda.is_available(): pixel_values = pixel_values.to(torch.bfloat16).cuda() with torch.inference_mode(): score = self.model(pixel_values).logits.squeeze().float().cpu().numpy() return score return AestheticPredictorV2_5_Impl() def load_anime_aesthetic_model(): """Load and return the Anime Aesthetic ONNX model.""" model_path = hf_hub_download(repo_id="skytnt/anime-aesthetic", filename="model.onnx") return rt.InferenceSession(model_path, providers=['CPUExecutionProvider']) def predict_anime_aesthetic(img, model): """Predict Anime Aesthetic score for a single image.""" img_np = np.array(img).astype(np.float32) / 255.0 s = 768 h, w = img_np.shape[:2] if h > w: new_h, new_w = s, int(s * w / h) else: new_h, new_w = int(s * h / w), s resized = cv2.resize(img_np, (new_w, new_h)) # Center the resized image in a square canvas canvas = np.zeros((s, s, 3), dtype=np.float32) pad_h = (s - new_h) // 2 pad_w = (s - new_w) // 2 canvas[pad_h:pad_h+new_h, pad_w:pad_w+new_w] = resized # Prepare input for model input_tensor = np.transpose(canvas, (2, 0, 1))[np.newaxis, :] pred = model.run(None, {"img": input_tensor})[0].item() return pred ##################################### # Image Evaluation Tool # ##################################### class ImageEvaluationTool: """Evaluation tool to process images through multiple aesthetic models and generate logs and HTML outputs.""" def __init__(self): self.device = 'cuda' if torch.cuda.is_available() else 'cpu' print(f"Using device: {self.device}") print("Loading models... This may take some time.") # Load models with progress logs print("Loading Aesthetic Shadow model...") self.aesthetic_shadow = pipeline("image-classification", model="NeoChen1024/aesthetic-shadow-v2-backup", device=self.device) print("Loading Waifu Scorer model...") self.waifu_scorer = WaifuScorer(device=self.device, verbose=True) print("Loading Aesthetic Predictor V2.5...") self.aesthetic_predictor = load_aesthetic_predictor_v2_5() print("Loading Anime Aesthetic model...") self.anime_aesthetic = load_anime_aesthetic_model() print("All models loaded successfully!") self.temp_dir = tempfile.mkdtemp() self.results = [] # Store final results for sorting and display self.available_models = { "aesthetic_shadow": {"name": "Aesthetic Shadow", "process": self._process_aesthetic_shadow}, "waifu_scorer": {"name": "Waifu Scorer", "process": self._process_waifu_scorer}, "aesthetic_predictor_v2_5": {"name": "Aesthetic V2.5", "process": self._process_aesthetic_predictor_v2_5}, "anime_aesthetic": {"name": "Anime Score", "process": self._process_anime_aesthetic}, } def image_to_base64(self, image: Image.Image) -> str: """Convert PIL Image to base64 encoded JPEG string.""" buffered = BytesIO() image.save(buffered, format="JPEG") return base64.b64encode(buffered.getvalue()).decode('utf-8') def auto_tune_batch_size(self, images: list) -> int: """Automatically determine the optimal batch size for processing.""" batch_size = 1 max_batch = len(images) test_image = images[0:1] while batch_size <= max_batch: try: if "aesthetic_shadow" in self.available_models and self.available_models["aesthetic_shadow"]['selected']: # Check if model is available and selected _ = self.aesthetic_shadow(test_image * batch_size) if "waifu_scorer" in self.available_models and self.available_models["waifu_scorer"]['selected']: # Check if model is available and selected _ = self.waifu_scorer(test_image * batch_size) if "aesthetic_predictor_v2_5" in self.available_models and self.available_models["aesthetic_predictor_v2_5"]['selected']: # Check if model is available and selected _ = self.aesthetic_predictor.inference(test_image * batch_size) batch_size *= 2 if batch_size > max_batch: break except Exception: break optimal = max(1, batch_size // 2) if optimal > 64: optimal = 64 print("Capped optimal batch size to 64") print(f"Optimal batch size determined: {optimal}") return optimal async def process_images_evaluation_with_logs(self, file_paths: list, auto_batch: bool, manual_batch_size: int, selected_models): """Asynchronously process images and yield updates with logs, HTML table, and progress bar.""" self.results = [] log_events = [] images = [] file_names = [] # Update available models based on selection for model_key in self.available_models: self.available_models[model_key]['selected'] = model_key in selected_models total_files = len(file_paths) log_events.append(f"Starting to load {total_files} images...") for f in file_paths: try: img = Image.open(f).convert("RGB") images.append(img) file_names.append(os.path.basename(f)) except Exception as e: log_events.append(f"Error opening {f}: {e}") if not images: log_events.append("No valid images loaded.") progress_percentage = 0 # Define progress_percentage here for no images case progress_html = self._generate_progress_html(progress_percentage) yield ("

No images loaded.

", "", self._format_logs(log_events), progress_html, manual_batch_size) return yield ("

Images loaded. Determining batch size...

", "", self._format_logs(log_events), self._generate_progress_html(0), manual_batch_size) await asyncio.sleep(0.1) try: manual_batch_size = int(manual_batch_size) if manual_batch_size is not None else 1 except ValueError: manual_batch_size = 1 log_events.append("Invalid manual batch size. Defaulting to 1.") optimal_batch = self.auto_tune_batch_size(images) if auto_batch else manual_batch_size log_events.append(f"Using batch size: {optimal_batch}") yield ("

Processing images in batches...

", "", self._format_logs(log_events), self._generate_progress_html(0), optimal_batch) await asyncio.sleep(0.1) total_images = len(images) for i in range(0, total_images, optimal_batch): batch_images = images[i:i+optimal_batch] batch_file_names = file_names[i:i+optimal_batch] batch_index = i // optimal_batch + 1 log_events.append(f"Processing batch {batch_index}: images {i+1} to {min(i+optimal_batch, total_images)}") batch_results = {} # Aesthetic Shadow processing if self.available_models['aesthetic_shadow']['selected']: batch_results['aesthetic_shadow'] = await self._process_aesthetic_shadow(batch_images, log_events) else: batch_results['aesthetic_shadow'] = [None] * len(batch_images) # Waifu Scorer processing if self.available_models['waifu_scorer']['selected']: batch_results['waifu_scorer'] = await self._process_waifu_scorer(batch_images, log_events) else: batch_results['waifu_scorer'] = [None] * len(batch_images) # Aesthetic Predictor V2.5 processing if self.available_models['aesthetic_predictor_v2_5']['selected']: batch_results['aesthetic_predictor_v2_5'] = await self._process_aesthetic_predictor_v2_5(batch_images, log_events) else: batch_results['aesthetic_predictor_v2_5'] = [None] * len(batch_images) # Anime Aesthetic processing (single image) if self.available_models['anime_aesthetic']['selected']: batch_results['anime_aesthetic'] = await self._process_anime_aesthetic(batch_images, log_events) else: batch_results['anime_aesthetic'] = [None] * len(batch_images) # Combine results for j in range(len(batch_images)): scores_to_average = [] for model_key in self.available_models: if self.available_models[model_key]['selected']: # Only consider selected models score = batch_results[model_key][j] if score is not None: scores_to_average.append(score) final_score = float(np.clip(np.mean(scores_to_average), 0.0, 10.0)) if scores_to_average else None thumbnail = batch_images[j].copy() thumbnail.thumbnail((200, 200)) result = { 'file_name': batch_file_names[j], 'img_data': self.image_to_base64(thumbnail), # Keep this for the HTML display 'final_score': final_score, } for model_key in self.available_models: # Add model scores to result if self.available_models[model_key]['selected']: result[model_key] = batch_results[model_key][j] self.results.append(result) self.sort_results() # Sort results after adding new result progress_percentage = min(100, ((i + len(batch_images)) / total_images) * 100) # Define progress_percentage here yield (f"

Processed batch {batch_index}.

", self.generate_html_table(self.results, selected_models), # Update table immediately self._format_logs(log_events[-10:]), self._generate_progress_html(progress_percentage), optimal_batch) await asyncio.sleep(0.1) log_events.append("All images processed.") self.sort_results() # Final sort after all images processed html_table = self.generate_html_table(self.results, selected_models) # Pass selected models to final table generation final_progress = self._generate_progress_html(100) yield ("

All images processed.

", html_table, self._format_logs(log_events[-10:]), final_progress, optimal_batch) async def _process_aesthetic_shadow(self, batch_images, log_events): try: shadow_results = self.aesthetic_shadow(batch_images) log_events.append("Aesthetic Shadow processed for batch.") except Exception as e: log_events.append(f"Error in Aesthetic Shadow: {e}") shadow_results = [None] * len(batch_images) aesthetic_shadow_scores = [] for res in shadow_results: try: hq_score = next(p for p in res if p['label'] == 'hq')['score'] score = float(np.clip(hq_score * 10.0, 0.0, 10.0)) except Exception: score = None aesthetic_shadow_scores.append(score) log_events.append("Aesthetic Shadow scores computed for batch.") return aesthetic_shadow_scores async def _process_waifu_scorer(self, batch_images, log_events): try: waifu_scores = self.waifu_scorer(batch_images) waifu_scores = [float(np.clip(s, 0.0, 10.0)) if s is not None else None for s in waifu_scores] log_events.append("Waifu Scorer processed for batch.") except Exception as e: log_events.append(f"Error in Waifu Scorer: {e}") waifu_scores = [None] * len(batch_images) return waifu_scores async def _process_aesthetic_predictor_v2_5(self, batch_images, log_events): try: v2_5_scores = self.aesthetic_predictor.inference(batch_images) v2_5_scores = [float(np.round(np.clip(s, 0.0, 10.0), 4)) if s is not None else None for s in v2_5_scores] log_events.append("Aesthetic Predictor V2.5 processed for batch.") except Exception as e: log_events.append(f"Error in Aesthetic Predictor V2.5: {e}") v2_5_scores = [None] * len(batch_images) return v2_5_scores async def _process_anime_aesthetic(self, batch_images, log_events): anime_scores = [] for j, img in enumerate(batch_images): try: score = predict_anime_aesthetic(img, self.anime_aesthetic) anime_scores.append(float(np.clip(score * 10.0, 0.0, 10.0))) log_events.append(f"Anime Aesthetic processed for image {j + 1}.") except Exception as e: log_events.append(f"Error in Anime Aesthetic for image {j + 1}: {e}") anime_scores.append(None) return anime_scores def _generate_progress_html(self, percentage: float) -> str: """Generate HTML for a progress bar given a percentage.""" return f"""
{percentage:.1f}%
""" def _format_logs(self, logs: list) -> str: """Format log events into an HTML string.""" return "
" + "
".join(logs) + "
" def sort_results(self, sort_by: str = "Final Score") -> list: """Sort results based on the specified column.""" key_map = { "Final Score": "final_score", "File Name": "file_name", "Aesthetic Shadow": "aesthetic_shadow", "Waifu Scorer": "waifu_scorer", "Aesthetic V2.5": "aesthetic_predictor_v2_5", "Anime Score": "anime_aesthetic" } key = key_map.get(sort_by, "final_score") reverse = sort_by != "File Name" self.results.sort(key=lambda r: r.get(key) if r.get(key) is not None else (-float('inf') if not reverse else float('inf')), reverse=reverse) return self.results def generate_html_table(self, results: list, selected_models) -> str: """Generate an HTML table to display the evaluation results.""" table_html = """ """ visible_models = [] # Keep track of visible model columns if "aesthetic_shadow" in selected_models: table_html += "" visible_models.append("aesthetic_shadow") if "waifu_scorer" in selected_models: table_html += "" visible_models.append("waifu_scorer") if "aesthetic_predictor_v2_5" in selected_models: table_html += "" visible_models.append("aesthetic_predictor_v2_5") if "anime_aesthetic" in selected_models: table_html += "" visible_models.append("anime_aesthetic") table_html += "" table_html += "" for result in results: table_html += "" table_html += f'' table_html += f'' for model_key in visible_models: # Iterate through visible models only score = result.get(model_key) table_html += self._format_score_cell(score) score = result.get("final_score") table_html += self._format_score_cell(score) table_html += "" table_html += """
Image File NameAesthetic ShadowWaifu ScorerAesthetic V2.5Anime ScoreFinal Score
{result["file_name"]}
""" return table_html def _format_score_cell(self, score): score_str = f"{score:.4f}" if isinstance(score, (int, float)) else "N/A" score_class = "" if isinstance(score, (int, float)): if score >= 7: score_class = "good-score" elif score >= 5: score_class = "medium-score" else: score_class = "bad-score" return f'{score_str}' def cleanup(self): """Clean up temporary directories.""" if os.path.exists(self.temp_dir): shutil.rmtree(self.temp_dir) ##################################### # Interface # ##################################### def create_interface(): evaluator = ImageEvaluationTool() sort_options = ["Final Score", "File Name", "Aesthetic Shadow", "Waifu Scorer", "Aesthetic V2.5", "Anime Score"] model_options = ["aesthetic_shadow", "waifu_scorer", "aesthetic_predictor_v2_5", "anime_aesthetic"] with gr.Blocks(theme=gr.themes.Soft()) as demo: gr.Markdown(""" # Comprehensive Image Evaluation Tool Upload images to evaluate them using multiple aesthetic and quality prediction models. **New features:** - **Dynamic Final Score:** Final score recalculates on model selection changes. - **Model Selection:** Choose which models to use for evaluation. - **Dynamic Table Updates:** Table updates automatically based on model selection. - **Automatic Sorting:** Table is automatically sorted by 'Final Score'. - **Detailed Logs:** See major processing events (limited to the last 10). - **Progress Bar:** Visual indication of processing status. - **Asynchronous Updates:** Streaming status and logs during processing. - **Batch Size Controls:** Choose manual batch size or let the tool auto-detect it. - **Download Results:** Export the evaluation results as CSV. """) with gr.Row(): with gr.Column(scale=1): input_images = gr.Files(label="Upload Images", file_count="multiple") model_checkboxes = gr.CheckboxGroup(model_options, label="Select Models", value=model_options, info="Choose models for evaluation.") auto_batch_checkbox = gr.Checkbox(label="Automatic Batch Size Detection", value=False, info="Enable to automatically determine the optimal batch size.") batch_size_input = gr.Number(label="Batch Size", value=1, interactive=True, info="Manually specify the batch size if auto-detection is disabled.") sort_dropdown = gr.Dropdown(sort_options, value="Final Score", label="Sort by", info="Select the column to sort results by.") process_btn = gr.Button("Evaluate Images", variant="primary") clear_btn = gr.Button("Clear Results") download_csv = gr.Button("Download CSV", variant="secondary") with gr.Column(scale=2): progress_bar = gr.HTML(label="Progress Bar", value="""
0%
""") log_window = gr.HTML(label="Detailed Logs", value="
Logs will appear here...
") status_html = gr.HTML(label="Status") output_html = gr.HTML(label="Evaluation Results") download_file_output = gr.File() # Initialize gr.File component without filename # Function to convert results to CSV format, excluding 'img_data'. def results_to_csv(selected_models): import csv import io if not evaluator.results: return None # Return None when no results are available output = io.StringIO() fieldnames = ['file_name', 'final_score'] # Base fieldnames for model_key in selected_models: # Add selected model names as fieldnames if model_key in selected_models: # Double check if model_key is indeed in selected_models list fieldnames.append(model_key) writer = csv.DictWriter(output, fieldnames=fieldnames) writer.writeheader() for res in evaluator.results: row_dict = {'file_name': res['file_name'], 'final_score': res['final_score']} # Base data for model_key in selected_models: # Add selected model scores if model_key in selected_models: # Double check before accessing res[model_key] row_dict[model_key] = res.get(model_key, 'N/A') # Use get with default 'N/A' if model not in result (shouldn't happen but for safety) writer.writerow(row_dict) return output.getvalue() def update_batch_size_interactivity(auto_batch): return gr.update(interactive=not auto_batch) async def process_images_and_update(files, auto_batch, manual_batch, selected_models): file_paths = [f.name for f in files] async for status, table, logs, progress, updated_batch in evaluator.process_images_evaluation_with_logs(file_paths, auto_batch, manual_batch, selected_models): yield status, table, logs, progress, gr.update(value=updated_batch, interactive=not auto_batch) def update_table_sort(sort_by_column, selected_models): sorted_results = evaluator.sort_results(sort_by_column) return evaluator.generate_html_table(sorted_results, selected_models) def update_table_model_selection(selected_models): # Recalculate final scores based on selected models for result in evaluator.results: scores_to_average = [] for model_key in evaluator.available_models: if model_key in selected_models and evaluator.available_models[model_key]['selected']: # consider only selected models from checkbox group and available_models score = result.get(model_key) if score is not None: scores_to_average.append(score) final_score = float(np.clip(np.mean(scores_to_average), 0.0, 10.0)) if scores_to_average else None result['final_score'] = final_score sorted_results = evaluator.sort_results() # Keep sorting by Final Score when models change return evaluator.generate_html_table(sorted_results, selected_models) def clear_results(): evaluator.results = [] return (gr.update(value=""), gr.update(value=""), gr.update(value=""), gr.update(value="""
0%
"""), gr.update(value=1)) def download_results_csv_trigger(selected_models): # Changed function name to avoid conflict and clarify purpose csv_content = results_to_csv(selected_models) if csv_content is None: return None # Indicate no file to download # Create a temporary file to save the CSV data with tempfile.NamedTemporaryFile(suffix=".csv", delete=False) as tmp_file: tmp_file.write(csv_content.encode()) temp_file_path = tmp_file.name # Get the path to the temporary file return temp_file_path # Return the path to the temporary file auto_batch_checkbox.change( update_batch_size_interactivity, inputs=[auto_batch_checkbox], outputs=[batch_size_input] ) process_btn.click( process_images_and_update, inputs=[input_images, auto_batch_checkbox, batch_size_input, model_checkboxes], outputs=[status_html, output_html, log_window, progress_bar, batch_size_input] ) sort_dropdown.change( update_table_sort, inputs=[sort_dropdown, model_checkboxes], outputs=[output_html] ) model_checkboxes.change( # Added change event for model checkboxes update_table_model_selection, inputs=[model_checkboxes], outputs=[output_html] ) clear_btn.click( clear_results, inputs=[], outputs=[status_html, output_html, log_window, progress_bar, batch_size_input] ) download_csv.click( download_results_csv_trigger, # Call the trigger function inputs=[model_checkboxes], outputs=[download_file_output] # Output is now the gr.File component ) demo.load(lambda: update_table_sort("Final Score", model_options), inputs=None, outputs=[output_html]) # Initial sort and table render gr.Markdown(""" ### Notes - Select models to use for evaluation using the checkboxes. - The 'Final Score' recalculates dynamically when models are selected/deselected. - The table updates automatically when models are selected/deselected and is always sorted by 'Final Score'. - The log window displays the most recent 10 events. - The progress bar shows overall processing status. - When 'Automatic Batch Size Detection' is enabled, the batch size field becomes disabled. - Use the download button to export your evaluation results as CSV. """) return demo if __name__ == "__main__": demo = create_interface() demo.queue().launch()