Update app.py

app.py

@@ -1,10 +1,9 @@
 import os
-import tempfile
-import base64
-from io import BytesIO
-from typing import List, Dict, Any, Optional, Tuple
-from dataclasses import dataclass
+import asyncio
+from typing import List, Dict, Optional, Tuple, Any
+from dataclasses import dataclass, field
 from pathlib import Path
+import logging
 
 import cv2
 import numpy as np
@@ -12,589 +11,508 @@ import torch
 import onnxruntime as rt
 from PIL import Image
 import gradio as gr
-import pandas as pd
 from transformers import pipeline
 from huggingface_hub import hf_hub_download
+import pandas as pd
+
+# Configure logging
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
+logger = logging.getLogger(__name__)
 
-# Import necessary function from aesthetic_predictor_v2_5
+# Import aesthetic predictor function
 from aesthetic_predictor_v2_5 import convert_v2_5_from_siglip
 
 
 @dataclass
 class EvaluationResult:
-    """Data class for storing image evaluation results."""
+    """Data class for storing image evaluation results"""
     file_name: str
-    image: Image.Image
-    aesthetic_shadow: Optional[float] = None
-    waifu_scorer: Optional[float] = None
-    aesthetic_v2_5: Optional[float] = None
-    anime_aesthetic: Optional[float] = None
+    image_path: str
+    scores: Dict[str, Optional[float]] = field(default_factory=dict)
     final_score: Optional[float] = None
+    
+    def calculate_final_score(self, selected_models: List[str]) -> None:
+        """Calculate the average score from selected models"""
+        valid_scores = [
+            score for model, score in self.scores.items() 
+            if model in selected_models and score is not None
+        ]
+        self.final_score = np.mean(valid_scores) if valid_scores else None
 
 
-class MLP(torch.nn.Module):
-    """Optimized MLP for image feature regression."""
-    def __init__(self, input_size: int = 768):
-        super().__init__()
-        self.network = torch.nn.Sequential(
-            torch.nn.Linear(input_size, 1024),
-            torch.nn.ReLU(),
-            torch.nn.BatchNorm1d(1024),
-            torch.nn.Dropout(0.2),
-            torch.nn.Linear(1024, 256),
-            torch.nn.ReLU(),
-            torch.nn.BatchNorm1d(256),
-            torch.nn.Dropout(0.1),
-            torch.nn.Linear(256, 64),
-            torch.nn.ReLU(),
-            torch.nn.Linear(64, 1)
-        )
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        return self.network(x)
+class BaseModel:
+    """Base class for all evaluation models"""
+    def __init__(self, name: str):
+        self.name = name
+        self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
+        
+    async def evaluate_batch(self, images: List[Image.Image]) -> List[Optional[float]]:
+        """Evaluate a batch of images"""
+        raise NotImplementedError
 
 
-class ModelLoader:
-    """Centralized model loading and management."""
-    
-    def __init__(self, device: str = None):
-        self.device = device or ('cuda' if torch.cuda.is_available() else 'cpu')
-        self.models = {}
-        self._load_all_models()
-    
-    def _load_all_models(self):
-        """Load all models during initialization."""
-        try:
-            self._load_aesthetic_shadow()
-            self._load_waifu_scorer()
-            self._load_aesthetic_v2_5()
-            self._load_anime_aesthetic()
-            print("✅ All models loaded successfully!")
-        except Exception as e:
-            print(f"❌ Error loading models: {e}")
-    
-    def _load_aesthetic_shadow(self):
-        """Load Aesthetic Shadow model."""
-        print("🔄 Loading Aesthetic Shadow...")
-        self.models['aesthetic_shadow'] = pipeline(
-            "image-classification",
+class AestheticShadowModel(BaseModel):
+    """Aesthetic Shadow V2 model implementation"""
+    def __init__(self):
+        super().__init__("Aesthetic Shadow")
+        logger.info(f"Loading {self.name} model...")
+        self.model = pipeline(
+            "image-classification", 
             model="NeoChen1024/aesthetic-shadow-v2-backup",
-            device=self.device
+            device=0 if self.device == 'cuda' else -1
         )
     
-    def _load_waifu_scorer(self):
-        """Load Waifu Scorer model."""
-        print("🔄 Loading Waifu Scorer...")
-        try:
-            import clip
-            
-            # Load MLP
-            model_path = hf_hub_download("Eugeoter/waifu-scorer-v3", "model.pth")
-            mlp = MLP()
-            state_dict = torch.load(model_path, map_location=self.device)
-            mlp.load_state_dict(state_dict)
-            mlp.to(self.device).eval()
-            
-            # Load CLIP
-            clip_model, preprocess = clip.load("ViT-L/14", device=self.device)
-            
-            self.models['waifu_scorer'] = {
-                'mlp': mlp,
-                'clip_model': clip_model,
-                'preprocess': preprocess
-            }
-        except Exception as e:
-            print(f"⚠️ Waifu Scorer not available: {e}")
-            self.models['waifu_scorer'] = None
-    
-    def _load_aesthetic_v2_5(self):
-        """Load Aesthetic Predictor V2.5."""
-        print("🔄 Loading Aesthetic V2.5...")
-        try:
-            model, preprocessor = convert_v2_5_from_siglip(
-                low_cpu_mem_usage=True,
-                trust_remote_code=True,
-            )
-            if torch.cuda.is_available():
-                model = model.to(torch.bfloat16).cuda()
-            
-            self.models['aesthetic_v2_5'] = {
-                'model': model,
-                'preprocessor': preprocessor
-            }
-        except Exception as e:
-            print(f"⚠️ Aesthetic V2.5 not available: {e}")
-            self.models['aesthetic_v2_5'] = None
-    
-    def _load_anime_aesthetic(self):
-        """Load Anime Aesthetic model."""
-        print("🔄 Loading Anime Aesthetic...")
+    async def evaluate_batch(self, images: List[Image.Image]) -> List[Optional[float]]:
         try:
-            model_path = hf_hub_download("skytnt/anime-aesthetic", "model.onnx")
-            self.models['anime_aesthetic'] = rt.InferenceSession(
-                model_path, 
-                providers=['CPUExecutionProvider']
-            )
+            results = self.model(images)
+            scores = []
+            for result in results:
+                hq_score = next((p['score'] for p in result if p['label'] == 'hq'), 0)
+                scores.append(float(np.clip(hq_score * 10.0, 0.0, 10.0)))
+            return scores
         except Exception as e:
-            print(f"⚠️ Anime Aesthetic not available: {e}")
-            self.models['anime_aesthetic'] = None
+            logger.error(f"Error in {self.name}: {e}")
+            return [None] * len(images)
 
 
-class ImageEvaluator:
-    """Main image evaluation class with batch processing."""
-    
+class WaifuScorerModel(BaseModel):
+    """Waifu Scorer V3 model implementation"""
     def __init__(self):
-        self.loader = ModelLoader()
-        self.temp_dir = Path(tempfile.mkdtemp())
+        super().__init__("Waifu Scorer")
+        logger.info(f"Loading {self.name} model...")
+        self._load_model()
     
-    def evaluate_images(
-        self,
-        images: List[Image.Image],
-        file_names: List[str],
-        selected_models: List[str],
-        batch_size: int = 4,
-        progress_callback=None
-    ) -> List[EvaluationResult]:
-        """Evaluate images using selected models."""
-        results = []
-        total_batches = (len(images) + batch_size - 1) // batch_size
-        
-        for batch_idx in range(0, len(images), batch_size):
-            batch_images = images[batch_idx:batch_idx + batch_size]
-            batch_names = file_names[batch_idx:batch_idx + batch_size]
+    def _load_model(self):
+        try:
+            import clip
             
-            # Update progress
-            if progress_callback:
-                progress = (batch_idx // batch_size + 1) / total_batches
-                progress_callback(progress, f"Processing batch {batch_idx//batch_size + 1}/{total_batches}")
+            # Load MLP model
+            self.mlp = self._create_mlp()
+            model_path = hf_hub_download("Eugeoter/waifu-scorer-v3", "model.pth")
+            state_dict = torch.load(model_path, map_location=self.device)
+            self.mlp.load_state_dict(state_dict)
+            self.mlp.to(self.device).eval()
             
-            # Process batch
-            batch_results = self._process_batch(batch_images, batch_names, selected_models)
-            results.extend(batch_results)
-        
-        return results
+            # Load CLIP model
+            self.clip_model, self.preprocess = clip.load("ViT-L/14", device=self.device)
+            self.available = True
+        except Exception as e:
+            logger.error(f"Failed to load {self.name}: {e}")
+            self.available = False
     
-    def _process_batch(
-        self,
-        images: List[Image.Image],
-        file_names: List[str],
-        selected_models: List[str]
-    ) -> List[EvaluationResult]:
-        """Process a single batch of images."""
-        batch_results = []
-        
-        # Initialize results
-        for i, (img, name) in enumerate(zip(images, file_names)):
-            result = EvaluationResult(file_name=name, image=img)
-            batch_results.append(result)
-        
-        # Process each selected model
-        if 'aesthetic_shadow' in selected_models:
-            scores = self._eval_aesthetic_shadow(images)
-            for result, score in zip(batch_results, scores):
-                result.aesthetic_shadow = score
-        
-        if 'waifu_scorer' in selected_models:
-            scores = self._eval_waifu_scorer(images)
-            for result, score in zip(batch_results, scores):
-                result.waifu_scorer = score
-        
-        if 'aesthetic_v2_5' in selected_models:
-            scores = self._eval_aesthetic_v2_5(images)
-            for result, score in zip(batch_results, scores):
-                result.aesthetic_v2_5 = score
-        
-        if 'anime_aesthetic' in selected_models:
-            scores = self._eval_anime_aesthetic(images)
-            for result, score in zip(batch_results, scores):
-                result.anime_aesthetic = score
-        
-        # Calculate final scores
-        for result in batch_results:
-            result.final_score = self._calculate_final_score(result, selected_models)
-        
-        return batch_results
+    def _create_mlp(self) -> torch.nn.Module:
+        """Create the MLP architecture"""
+        return torch.nn.Sequential(
+            torch.nn.Linear(768, 2048),
+            torch.nn.ReLU(),
+            torch.nn.BatchNorm1d(2048),
+            torch.nn.Dropout(0.3),
+            torch.nn.Linear(2048, 512),
+            torch.nn.ReLU(),
+            torch.nn.BatchNorm1d(512),
+            torch.nn.Dropout(0.3),
+            torch.nn.Linear(512, 256),
+            torch.nn.ReLU(),
+            torch.nn.BatchNorm1d(256),
+            torch.nn.Dropout(0.2),
+            torch.nn.Linear(256, 128),
+            torch.nn.ReLU(),
+            torch.nn.BatchNorm1d(128),
+            torch.nn.Dropout(0.1),
+            torch.nn.Linear(128, 32),
+            torch.nn.ReLU(),
+            torch.nn.Linear(32, 1)
+        )
     
-    def _eval_aesthetic_shadow(self, images: List[Image.Image]) -> List[Optional[float]]:
-        """Evaluate using Aesthetic Shadow model."""
-        if not self.loader.models.get('aesthetic_shadow'):
+    @torch.no_grad()
+    async def evaluate_batch(self, images: List[Image.Image]) -> List[Optional[float]]:
+        if not self.available:
             return [None] * len(images)
         
         try:
-            results = self.loader.models['aesthetic_shadow'](images)
-            scores = []
-            for result in results:
-                try:
-                    hq_score = next(p for p in result if p['label'] == 'hq')['score']
-                    scores.append(float(np.clip(hq_score * 10.0, 0.0, 10.0)))
-                except:
-                    scores.append(None)
+            # Process images
+            image_tensors = torch.cat([self.preprocess(img).unsqueeze(0) for img in images])
+            image_tensors = image_tensors.to(self.device)
+            
+            # Extract features and predict
+            features = self.clip_model.encode_image(image_tensors)
+            features = features / features.norm(dim=-1, keepdim=True)
+            predictions = self.mlp(features)
+            
+            scores = predictions.clamp(0, 10).cpu().numpy().flatten().tolist()
             return scores
         except Exception as e:
-            print(f"Error in Aesthetic Shadow: {e}")
-            return [None] * len(images)
-    
-    def _eval_waifu_scorer(self, images: List[Image.Image]) -> List[Optional[float]]:
-        """Evaluate using Waifu Scorer model."""
-        model_dict = self.loader.models.get('waifu_scorer')
-        if not model_dict:
-            return [None] * len(images)
-        
-        try:
-            with torch.no_grad():
-                # Preprocess images
-                image_tensors = [model_dict['preprocess'](img).unsqueeze(0) for img in images]
-                if len(image_tensors) == 1:
-                    image_tensors = image_tensors * 2  # CLIP requirement
-                
-                image_batch = torch.cat(image_tensors).to(self.loader.device)
-                image_features = model_dict['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(self.loader.device)
-                
-                predictions = model_dict['mlp'](im_emb)
-                scores = predictions.clamp(0, 10).cpu().numpy().flatten().tolist()
-                
-                return scores[:len(images)]
-        except Exception as e:
-            print(f"Error in Waifu Scorer: {e}")
+            logger.error(f"Error in {self.name}: {e}")
             return [None] * len(images)
+
+
+class AestheticPredictorV25Model(BaseModel):
+    """Aesthetic Predictor V2.5 model implementation"""
+    def __init__(self):
+        super().__init__("Aesthetic V2.5")
+        logger.info(f"Loading {self.name} model...")
+        self.model, self.preprocessor = convert_v2_5_from_siglip(
+            low_cpu_mem_usage=True,
+            trust_remote_code=True,
+        )
+        if self.device == 'cuda':
+            self.model = self.model.to(torch.bfloat16).cuda()
     
-    def _eval_aesthetic_v2_5(self, images: List[Image.Image]) -> List[Optional[float]]:
-        """Evaluate using Aesthetic Predictor V2.5."""
-        model_dict = self.loader.models.get('aesthetic_v2_5')
-        if not model_dict:
-            return [None] * len(images)
-        
+    @torch.no_grad()
+    async def evaluate_batch(self, images: List[Image.Image]) -> List[Optional[float]]:
         try:
-            rgb_images = [img.convert("RGB") for img in images]
-            pixel_values = model_dict['preprocessor'](images=rgb_images, return_tensors="pt").pixel_values
+            images_rgb = [img.convert("RGB") for img in images]
+            pixel_values = self.preprocessor(images=images_rgb, return_tensors="pt").pixel_values
             
-            if torch.cuda.is_available():
+            if self.device == 'cuda':
                 pixel_values = pixel_values.to(torch.bfloat16).cuda()
             
-            with torch.inference_mode():
-                scores = model_dict['model'](pixel_values).logits.squeeze().float().cpu().numpy()
-                if scores.ndim == 0:
-                    scores = np.array([scores])
-                
-                return [float(np.clip(s, 0.0, 10.0)) for s in scores.tolist()]
+            scores = self.model(pixel_values).logits.squeeze().float().cpu().numpy()
+            if scores.ndim == 0:
+                scores = np.array([scores])
+            
+            return [float(np.clip(s, 0.0, 10.0)) for s in scores]
         except Exception as e:
-            print(f"Error in Aesthetic V2.5: {e}")
+            logger.error(f"Error in {self.name}: {e}")
             return [None] * len(images)
+
+
+class AnimeAestheticModel(BaseModel):
+    """Anime Aesthetic model implementation"""
+    def __init__(self):
+        super().__init__("Anime Score")
+        logger.info(f"Loading {self.name} model...")
+        model_path = hf_hub_download(repo_id="skytnt/anime-aesthetic", filename="model.onnx")
+        self.session = rt.InferenceSession(model_path, providers=['CPUExecutionProvider'])
     
-    def _eval_anime_aesthetic(self, images: List[Image.Image]) -> List[Optional[float]]:
-        """Evaluate using Anime Aesthetic model."""
-        model = self.loader.models.get('anime_aesthetic')
-        if not model:
-            return [None] * len(images)
-        
+    async def evaluate_batch(self, images: List[Image.Image]) -> List[Optional[float]]:
         scores = []
         for img in images:
             try:
-                # Preprocess image
-                img_np = np.array(img).astype(np.float32) / 255.0
-                h, w = img_np.shape[:2]
-                s = 768
-                
-                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))
-                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
-                
-                input_tensor = np.transpose(canvas, (2, 0, 1))[np.newaxis, :]
-                pred = model.run(None, {"img": input_tensor})[0].item()
-                scores.append(float(np.clip(pred * 10.0, 0.0, 10.0)))
+                score = self._process_single_image(img)
+                scores.append(float(np.clip(score * 10.0, 0.0, 10.0)))
             except Exception as e:
-                print(f"Error processing image: {e}")
+                logger.error(f"Error in {self.name} for single image: {e}")
                 scores.append(None)
-        
         return scores
     
-    def _calculate_final_score(self, result: EvaluationResult, selected_models: List[str]) -> Optional[float]:
-        """Calculate final score from selected model results."""
-        scores = []
+    def _process_single_image(self, img: Image.Image) -> float:
+        """Process a single image through the model"""
+        img_np = np.array(img).astype(np.float32) / 255.0
+        size = 768
+        h, w = img_np.shape[:2]
+        
+        # Calculate new dimensions
+        if h > w:
+            new_h, new_w = size, int(size * w / h)
+        else:
+            new_h, new_w = int(size * h / w), size
+        
+        # Resize and center
+        resized = cv2.resize(img_np, (new_w, new_h))
+        canvas = np.zeros((size, size, 3), dtype=np.float32)
+        pad_h = (size - new_h) // 2
+        pad_w = (size - new_w) // 2
+        canvas[pad_h:pad_h+new_h, pad_w:pad_w+new_w] = resized
+        
+        # Prepare input
+        input_tensor = np.transpose(canvas, (2, 0, 1))[np.newaxis, :]
+        return self.session.run(None, {"img": input_tensor})[0].item()
+
+
+class ImageEvaluator:
+    """Main class for managing image evaluation"""
+    def __init__(self):
+        self.models: Dict[str, BaseModel] = {}
+        self._initialize_models()
+        self.results: List[EvaluationResult] = []
+        
+    def _initialize_models(self):
+        """Initialize all evaluation models"""
+        model_classes = [
+            ("aesthetic_shadow", AestheticShadowModel),
+            ("waifu_scorer", WaifuScorerModel),
+            ("aesthetic_predictor_v2_5", AestheticPredictorV25Model),
+            ("anime_aesthetic", AnimeAestheticModel),
+        ]
+        
+        for key, model_class in model_classes:
+            try:
+                self.models[key] = model_class()
+                logger.info(f"Successfully loaded {key}")
+            except Exception as e:
+                logger.error(f"Failed to load {key}: {e}")
+    
+    async def evaluate_images(
+        self, 
+        file_paths: List[str], 
+        selected_models: List[str],
+        batch_size: int = 8,
+        progress_callback = None
+    ) -> Tuple[List[EvaluationResult], List[str]]:
+        """Evaluate images with selected models"""
+        logs = []
+        results = []
+        
+        # Load images
+        images = []
+        valid_paths = []
+        for path in file_paths:
+            try:
+                img = Image.open(path).convert("RGB")
+                images.append(img)
+                valid_paths.append(path)
+            except Exception as e:
+                logs.append(f"Failed to load {Path(path).name}: {e}")
+        
+        if not images:
+            logs.append("No valid images to process")
+            return results, logs
+        
+        logs.append(f"Loaded {len(images)} images")
         
-        for model in selected_models:
-            score = getattr(result, model, None)
-            if score is not None:
-                scores.append(score)
+        # Process in batches
+        total_batches = (len(images) + batch_size - 1) // batch_size
         
-        return float(np.mean(scores)) if scores else None
+        for batch_idx in range(0, len(images), batch_size):
+            batch_images = images[batch_idx:batch_idx + batch_size]
+            batch_paths = valid_paths[batch_idx:batch_idx + batch_size]
+            
+            # Evaluate with each selected model
+            batch_results = {}
+            for model_key in selected_models:
+                if model_key in self.models:
+                    scores = await self.models[model_key].evaluate_batch(batch_images)
+                    batch_results[model_key] = scores
+                    logs.append(f"Processed batch {batch_idx//batch_size + 1}/{total_batches} with {self.models[model_key].name}")
+            
+            # Create results
+            for i, (path, img) in enumerate(zip(batch_paths, batch_images)):
+                result = EvaluationResult(
+                    file_name=Path(path).name,
+                    image_path=path
+                )
+                
+                for model_key in selected_models:
+                    if model_key in batch_results:
+                        result.scores[model_key] = batch_results[model_key][i]
+                
+                result.calculate_final_score(selected_models)
+                results.append(result)
+            
+            # Update progress
+            if progress_callback:
+                progress = (batch_idx + batch_size) / len(images) * 100
+                progress_callback(min(progress, 100))
+        
+        self.results = results
+        return results, logs
     
-    def results_to_dataframe(self, results: List[EvaluationResult]) -> pd.DataFrame:
-        """Convert results to pandas DataFrame."""
+    def get_results_dataframe(self, selected_models: List[str]) -> pd.DataFrame:
+        """Convert results to pandas DataFrame"""
+        if not self.results:
+            return pd.DataFrame()
+        
         data = []
-        for result in results:
+        for result in self.results:
             row = {
                 'File Name': result.file_name,
-                'Final Score': result.final_score,
+                'Image': result.image_path,
             }
-            if result.aesthetic_shadow is not None:
-                row['Aesthetic Shadow'] = result.aesthetic_shadow
-            if result.waifu_scorer is not None:
-                row['Waifu Scorer'] = result.waifu_scorer
-            if result.aesthetic_v2_5 is not None:
-                row['Aesthetic V2.5'] = result.aesthetic_v2_5
-            if result.anime_aesthetic is not None:
-                row['Anime Aesthetic'] = result.anime_aesthetic
+            
+            # Add model scores
+            for model_key in selected_models:
+                if model_key in self.models:
+                    score = result.scores.get(model_key)
+                    row[self.models[model_key].name] = f"{score:.4f}" if score is not None else "N/A"
+            
+            row['Final Score'] = f"{result.final_score:.4f}" if result.final_score is not None else "N/A"
             data.append(row)
         
         return pd.DataFrame(data)
-    
-    def optimize_batch_size(self, sample_images: List[Image.Image]) -> int:
-        """Automatically determine optimal batch size."""
-        if not sample_images:
-            return 1
-        
-        test_image = sample_images[0]
-        batch_size = 1
-        max_test = min(16, len(sample_images))
-        
-        while batch_size <= max_test:
-            try:
-                test_batch = [test_image] * batch_size
-                # Test with a lightweight model
-                if self.loader.models.get('aesthetic_shadow'):
-                    _ = self.loader.models['aesthetic_shadow'](test_batch)
-                batch_size *= 2
-            except Exception:
-                break
-        
-        optimal = max(1, batch_size // 2)
-        return min(optimal, 8)  # Cap at reasonable size
 
 
 def create_interface():
-    """Create the Gradio interface."""
+    """Create the Gradio interface"""
     evaluator = ImageEvaluator()
     
-    # Available models
-    model_choices = [
+    # Model options for checkbox
+    model_options = [
         ("Aesthetic Shadow", "aesthetic_shadow"),
         ("Waifu Scorer", "waifu_scorer"),
-        ("Aesthetic V2.5", "aesthetic_v2_5"),
-        ("Anime Aesthetic", "anime_aesthetic")
+        ("Aesthetic V2.5", "aesthetic_predictor_v2_5"),
+        ("Anime Score", "anime_aesthetic")
     ]
-    available_models = [choice[1] for choice in model_choices]
     
-    with gr.Blocks(title="Image Evaluation Tool", theme=gr.themes.Soft()) as app:
+    with gr.Blocks(theme=gr.themes.Soft(), title="Image Evaluation Tool") as demo:
         gr.Markdown("""
-        # 🎨 Modern Image Evaluation Tool
+        # 🎨 Advanced Image Evaluation Tool
         
-        Upload images to evaluate them using state-of-the-art aesthetic and quality prediction models.
-        
-        **Features:**
-        - Multiple AI models for comprehensive evaluation
-        - Batch processing with automatic optimization
-        - Interactive results table with sorting and filtering
-        - CSV export functionality
-        - Real-time progress tracking
+        Evaluate images using state-of-the-art aesthetic and quality prediction models.
+        Upload your images and select the models you want to use for evaluation.
         """)
         
         with gr.Row():
             with gr.Column(scale=1):
-                # Input components
                 input_files = gr.File(
-                    label="📁 Upload Images",
+                    label="Upload Images",
                     file_count="multiple",
                     file_types=["image"]
                 )
                 
-                model_selection = gr.CheckboxGroup(
-                    choices=model_choices,
-                    value=available_models,
-                    label="🤖 Select Models",
+                model_checkboxes = gr.CheckboxGroup(
+                    choices=[label for label, _ in model_options],
+                    value=[label for label, _ in model_options],
+                    label="Select Models",
                     info="Choose which models to use for evaluation"
                 )
                 
                 with gr.Row():
-                    auto_batch = gr.Checkbox(
-                        label="🔄 Auto Batch Size",
-                        value=True,
-                        info="Automatically optimize batch size"
-                    )
-                    
-                    manual_batch = gr.Slider(
+                    batch_size = gr.Slider(
                         minimum=1,
-                        maximum=16,
-                        value=4,
+                        maximum=64,
+                        value=8,
                         step=1,
-                        label="📊 Batch Size",
-                        interactive=False,
-                        info="Manual batch size (when auto is disabled)"
+                        label="Batch Size",
+                        info="Number of images to process at once"
                     )
                 
-                evaluate_btn = gr.Button(
-                    "🚀 Evaluate Images",
-                    variant="primary",
-                    size="lg"
-                )
-                
-                clear_btn = gr.Button("🗑️ Clear Results", variant="secondary")
+                with gr.Row():
+                    evaluate_btn = gr.Button("🚀 Evaluate Images", variant="primary", scale=2)
+                    clear_btn = gr.Button("🗑️ Clear", variant="secondary", scale=1)
             
             with gr.Column(scale=2):
-                # Progress and status
-                progress_bar = gr.Progress()
-                status_text = gr.Textbox(
-                    label="📊 Status",
-                    interactive=False,
-                    max_lines=2
+                progress = gr.Progress()
+                logs = gr.Textbox(
+                    label="Processing Logs",
+                    lines=10,
+                    max_lines=10,
+                    autoscroll=True
                 )
                 
-                # Results display
-                results_table = gr.DataFrame(
-                    label="📋 Evaluation Results",
+                results_df = gr.Dataframe(
+                    label="Evaluation Results",
                     interactive=False,
-                    wrap=True,
-                    max_height=400
+                    wrap=True
                 )
                 
-                # Export functionality
-                with gr.Row():
-                    export_csv = gr.Button("📥 Export CSV", variant="secondary")
-                    download_file = gr.File(
-                        label="💾 Download",
-                        visible=False
-                    )
+                download_btn = gr.Button("📥 Download Results (CSV)", variant="secondary")
+                download_file = gr.File(visible=False)
         
-        # State management
+        # State for storing results
         results_state = gr.State([])
         
-        # Event handlers
-        def toggle_batch_slider(auto_enabled):
-            return gr.update(interactive=not auto_enabled)
-        
-        def process_images(files, models, auto_batch_enabled, manual_batch_size, progress=gr.Progress()):
-            if not files or not models:
-                return "❌ Please upload images and select at least one model", pd.DataFrame(), []
+        async def process_images(files, selected_model_labels, batch_size, progress=gr.Progress()):
+            """Process uploaded images"""
+            if not files:
+                return "Please upload images first", pd.DataFrame(), []
             
-            try:
-                # Load images
-                images = []
-                file_names = []
-                
-                progress(0.1, "📂 Loading images...")
-                
-                for file in files:
-                    try:
-                        img = Image.open(file.name).convert("RGB")
-                        images.append(img)
-                        file_names.append(os.path.basename(file.name))
-                    except Exception as e:
-                        print(f"Error loading {file.name}: {e}")
-                
-                if not images:
-                    return "❌ No valid images loaded", pd.DataFrame(), []
-                
-                # Determine batch size
-                if auto_batch_enabled:
-                    batch_size = evaluator.optimize_batch_size(images[:2])
-                    progress(0.2, f"🔧 Optimized batch size: {batch_size}")
-                else:
-                    batch_size = int(manual_batch_size)
-                
-                # Process images
-                def progress_callback(prog, msg):
-                    progress(0.2 + prog * 0.7, msg)
-                
-                results = evaluator.evaluate_images(
-                    images, file_names, models, batch_size, progress_callback
-                )
-                
-                progress(0.95, "📊 Generating results table...")
-                
-                # Convert to DataFrame
-                df = evaluator.results_to_dataframe(results)
-                df = df.sort_values('Final Score', ascending=False, na_position='last')
-                
-                progress(1.0, f"✅ Processed {len(results)} images successfully!")
-                
-                return f"✅ Evaluated {len(results)} images using {len(models)} models", df, results
-                
-            except Exception as e:
-                return f"❌ Error during processing: {str(e)}", pd.DataFrame(), []
+            # Convert labels to keys
+            selected_models = [key for label, key in model_options if label in selected_model_labels]
+            
+            # Get file paths
+            file_paths = [f.name for f in files]
+            
+            # Progress callback
+            def update_progress(value):
+                progress(value / 100, desc=f"Processing images... {value:.0f}%")
+            
+            # Evaluate images
+            results, logs = await evaluator.evaluate_images(
+                file_paths,
+                selected_models,
+                batch_size,
+                update_progress
+            )
+            
+            # Create DataFrame
+            df = evaluator.get_results_dataframe(selected_models)
+            
+            # Format logs
+            log_text = "\n".join(logs[-10:])  # Show last 10 logs
+            
+            return log_text, df, results
         
-        def update_results_table(models, current_results):
-            if not current_results:
+        def update_results_on_model_change(selected_model_labels, results):
+            """Update results when model selection changes"""
+            if not results:
                 return pd.DataFrame()
             
-            # Recalculate final scores based on selected models
-            for result in current_results:
-                result.final_score = evaluator._calculate_final_score(result, models)
+            # Convert labels to keys
+            selected_models = [key for label, key in model_options if label in selected_model_labels]
             
-            df = evaluator.results_to_dataframe(current_results)
-            return df.sort_values('Final Score', ascending=False, na_position='last')
-        
-        def export_results(current_results):
-            if not current_results:
-                return gr.update(visible=False)
+            # Recalculate final scores
+            for result in results:
+                result.calculate_final_score(selected_models)
             
-            df = evaluator.results_to_dataframe(current_results)
-            csv_path = evaluator.temp_dir / "evaluation_results.csv"
-            df.to_csv(csv_path, index=False)
+            # Update evaluator results
+            evaluator.results = results
             
-            return gr.update(value=str(csv_path), visible=True)
+            # Create updated DataFrame
+            return evaluator.get_results_dataframe(selected_models)
         
-        def clear_all():
-            return (
-                "🔄 Ready for new evaluation",
-                pd.DataFrame(),
-                [],
-                gr.update(visible=False)
-            )
+        def clear_interface():
+            """Clear all results"""
+            return "", pd.DataFrame(), [], None
         
-        # Wire up events
-        auto_batch.change(
-            toggle_batch_slider,
-            inputs=[auto_batch],
-            outputs=[manual_batch]
-        )
+        def prepare_download(selected_model_labels, results):
+            """Prepare CSV file for download"""
+            if not results:
+                return None
+            
+            # Convert labels to keys
+            selected_models = [key for label, key in model_options if label in selected_model_labels]
+            
+            # Get DataFrame
+            df = evaluator.get_results_dataframe(selected_models)
+            
+            # Save to temporary file
+            import tempfile
+            with tempfile.NamedTemporaryFile(mode='w', suffix='.csv', delete=False) as f:
+                df.to_csv(f, index=False)
+                return f.name
         
+        # Event handlers
         evaluate_btn.click(
-            process_images,
-            inputs=[input_files, model_selection, auto_batch, manual_batch],
-            outputs=[status_text, results_table, results_state]
+            fn=process_images,
+            inputs=[input_files, model_checkboxes, batch_size],
+            outputs=[logs, results_df, results_state]
         )
         
-        model_selection.change(
-            update_results_table,
-            inputs=[model_selection, results_state],
-            outputs=[results_table]
+        model_checkboxes.change(
+            fn=update_results_on_model_change,
+            inputs=[model_checkboxes, results_state],
+            outputs=[results_df]
         )
         
-        export_csv.click(
-            export_results,
-            inputs=[results_state],
-            outputs=[download_file]
+        clear_btn.click(
+            fn=clear_interface,
+            outputs=[logs, results_df, results_state, download_file]
         )
         
-        clear_btn.click(
-            clear_all,
-            outputs=[status_text, results_table, results_state, download_file]
+        download_btn.click(
+            fn=prepare_download,
+            inputs=[model_checkboxes, results_state],
+            outputs=[download_file]
         )
         
-        # Initial setup
-        app.load(lambda: "🔄 Ready for evaluation - Upload images to get started!")
+        gr.Markdown("""
+        ### 📝 Notes
+        - **Model Selection**: Choose which models to use for evaluation. Final score is the average of selected models.
+        - **Batch Size**: Adjust based on your GPU memory. Larger batches process faster.
+        - **Results Table**: Click column headers to sort. The table updates automatically when models are selected/deselected.
+        - **Download**: Export results as CSV for further analysis.
+        
+        ### 🎯 Score Interpretation
+        - **7-10**: High quality/aesthetic appeal
+        - **5-7**: Medium quality
+        - **0-5**: Lower quality
+        """)
     
-    return app
+    return demo
 
 
 if __name__ == "__main__":
-    app = create_interface()
-    app.queue(max_size=10).launch(
-        server_name="0.0.0.0",
-        server_port=7860,
-        share=False,
-        show_error=True
-    )
+    # Create and launch the interface
+    demo = create_interface()
+    demo.queue().launch()