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 ModelManager: """Manages model loading and processing requests using a queue.""" 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 once during initialization print("Loading Aesthetic Shadow model...") self.aesthetic_shadow_model = pipeline("image-classification", model="NeoChen1024/aesthetic-shadow-v2-backup", device=self.device) print("Loading Waifu Scorer model...") self.waifu_scorer_model = WaifuScorer(device=self.device, verbose=True) print("Loading Aesthetic Predictor V2.5...") self.aesthetic_predictor_model = load_aesthetic_predictor_v2_5() print("Loading Anime Aesthetic model...") self.anime_aesthetic_model = load_anime_aesthetic_model() print("All models loaded successfully!") self.available_models = { "aesthetic_shadow": {"name": "Aesthetic Shadow", "process": self._process_aesthetic_shadow, "model": self.aesthetic_shadow_model}, "waifu_scorer": {"name": "Waifu Scorer", "process": self._process_waifu_scorer, "model": self.waifu_scorer_model}, "aesthetic_predictor_v2_5": {"name": "Aesthetic V2.5", "process": self._process_aesthetic_predictor_v2_5, "model": self.aesthetic_predictor_model}, "anime_aesthetic": {"name": "Anime Score", "process": self._process_anime_aesthetic, "model": self.anime_aesthetic_model}, } self.processing_queue: asyncio.Queue = asyncio.Queue() self.worker_task = None # Initialize worker_task to None self.temp_dir = tempfile.mkdtemp() async def start_worker(self): """Start the background worker task.""" if self.worker_task is None: self.worker_task = asyncio.create_task(self._worker()) async def _worker(self): """Background worker to process image evaluation requests from the queue.""" while True: request = await self.processing_queue.get() if request is None: # Shutdown signal self.processing_queue.task_done() break try: results = await self._process_request(request) request['results_future'].set_result(results) # Fulfill the future with results except Exception as e: request['results_future'].set_exception(e) # Set exception if processing fails finally: self.processing_queue.task_done() async def submit_request(self, request_data): """Submit a new image processing request to the queue.""" results_future = asyncio.Future() # Future to hold the results request = {**request_data, 'results_future': results_future} await self.processing_queue.put(request) return await results_future # Wait for and return results async def _process_request(self, request): """Process a single image evaluation request.""" file_paths = request['file_paths'] auto_batch = request['auto_batch'] manual_batch_size = request['manual_batch_size'] selected_models = request['selected_models'] log_events = [] images = [] file_names = [] final_results = [] # Prepare images and file names 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.") return [], log_events, 0, manual_batch_size log_events.append("Images loaded. Determining batch size...") 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}") 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 = {} # Process selected models for model_key in selected_models: if self.available_models[model_key]['selected']: # Ensure model is selected batch_results[model_key] = await self.available_models[model_key]['process'](batch_images, log_events) # Removed 'self' here else: batch_results[model_key] = [None] * len(batch_images) # Combine results and create final results list for j in range(len(batch_images)): scores_to_average = [] for model_key in selected_models: if self.available_models[model_key]['selected']: # Ensure model is selected 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 selected_models: # Add model scores to result if self.available_models[model_key]['selected']: result[model_key] = batch_results[model_key][j] final_results.append(result) log_events.append("All images processed.") return final_results, log_events, 100, optimal_batch 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.available_models["aesthetic_shadow"]['model'](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.available_models["waifu_scorer"]['model'](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.available_models["aesthetic_predictor_v2_5"]['model'].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(f"Optimal batch size determined: {optimal}") print(f"Optimal batch size determined: {optimal}") return optimal async def _process_aesthetic_shadow(self, batch_images, log_events): try: shadow_results = self.available_models["aesthetic_shadow"]['model'](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.available_models["waifu_scorer"]['model'](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.available_models["aesthetic_predictor_v2_5"]['model'].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.available_models["anime_aesthetic"]['model']) 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, results, 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" 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 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 and shutdown worker.""" if os.path.exists(self.temp_dir): shutil.rmtree(self.temp_dir) if self.worker_task is not None: # Check if worker_task was started asyncio.run(self.shutdown()) # Shutdown worker gracefully async def shutdown(self): """Send shutdown signal to worker and wait for it to finish.""" if self.worker_task is not None: # Check if worker_task was started await self.processing_queue.put(None) # Send shutdown signal await self.worker_task # Wait for worker task to complete await self.processing_queue.join() # Wait for queue to be empty ##################################### # Interface # ##################################### model_manager = ModelManager() # Initialize ModelManager once outside the interface function def create_interface(): 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 global_results_state = gr.State([]) # Initialize a global state to hold results # Function to convert results to CSV format, excluding 'img_data'. def results_to_csv(results, selected_models): # Take results as input import csv import io if not 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 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, current_results): file_paths = [f.name for f in files] # Prepare request data for the ModelManager request_data = { 'file_paths': file_paths, 'auto_batch': auto_batch, 'manual_batch_size': manual_batch, 'selected_models': {model: {'selected': model in selected_models} for model in model_options} # Pass model selections } # Submit request and get results from ModelManager results, logs, progress_percent, updated_batch = await model_manager.submit_request(request_data) updated_results = current_results + results # Append new results to current results html_table = model_manager.generate_html_table(updated_results, selected_models) progress_html = model_manager._generate_progress_html(progress_percent) log_html = model_manager._format_logs(logs[-10:]) return status_html, html_table, log_html, progress_html, gr.update(value=updated_batch, interactive=not auto_batch), updated_results def update_table_sort(sort_by_column, selected_models, current_results): sorted_results = model_manager.sort_results(current_results, sort_by_column) return model_manager.generate_html_table(sorted_results, selected_models), sorted_results # Return sorted results def update_table_model_selection(selected_models, current_results): # Recalculate final scores based on selected models for result in current_results: scores_to_average = [] for model_key in model_options: # Use model_options here, not available_models from manager in UI context if model_key in selected_models and model_key in model_manager.available_models and model_manager.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 = model_manager.sort_results(current_results, "Final Score") # Keep sorting by Final Score when models change return model_manager.generate_html_table(sorted_results, selected_models), sorted_results def clear_results(): return (gr.update(value=""), gr.update(value=""), gr.update(value=""), gr.update(value="""
0%
"""), gr.update(value=1), []) # Clear results state def download_results_csv_trigger(selected_models, current_results): # Changed function name to avoid conflict and clarify purpose csv_content = results_to_csv(current_results, 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 # Set initial selection state for models in ModelManager (important!) for model_key in model_options: model_manager.available_models[model_key]['selected'] = True # Default to all selected initially 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, global_results_state], outputs=[status_html, output_html, log_window, progress_bar, batch_size_input, global_results_state] ) sort_dropdown.change( update_table_sort, inputs=[sort_dropdown, model_checkboxes, global_results_state], outputs=[output_html, global_results_state] ) model_checkboxes.change( # Added change event for model checkboxes update_table_model_selection, inputs=[model_checkboxes, global_results_state], outputs=[output_html, global_results_state] ) clear_btn.click( clear_results, inputs=[], outputs=[status_html, output_html, log_window, progress_bar, batch_size_input, global_results_state] ) download_csv.click( download_results_csv_trigger, # Call the trigger function inputs=[model_checkboxes, global_results_state], 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, global_results_state]) # Initial sort and table render, pass empty initial results demo.load(model_manager.start_worker) # Start the worker task on demo load 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()