Update app.py

app.py

@@ -10,8 +10,12 @@ from huggingface_hub import hf_hub_download
 import pandas as pd
 import tempfile
 import shutil
+import base64
+from io import BytesIO
+
+# Import necessary function from aesthetic_predictor_v2_5
+from aesthetic_predictor_v2_5 import convert_v2_5_from_siglip
 
-# Utility classes and functions from provided code
 class MLP(torch.nn.Module):
     def __init__(self, input_size, xcol='emb', ycol='avg_rating', batch_norm=True):
         super().__init__()
@@ -43,90 +47,100 @@ class MLP(torch.nn.Module):
     def forward(self, x):
         return self.layers(x)
 
-
 class WaifuScorer(object):
     def __init__(self, model_path=None, device='cuda', cache_dir=None, verbose=False):
         self.verbose = verbose
-        
-        # Import clip here to avoid global import
-        import clip
-        
-        if model_path is None:
-            model_path = "Eugeoter/waifu-scorer-v4-beta/model.pth"
-            if self.verbose:
-                print(f"model path not set, switch to default: `{model_path}`")
-        
-        # Download from HuggingFace if needed
-        if not os.path.isfile(model_path):
-            split = model_path.split("/")
-            username, repo_id, model_name = split[-3], split[-2], split[-1]
-            model_path = hf_hub_download(f"{username}/{repo_id}", model_name, cache_dir=cache_dir)
-            
-        print(f"Loading WaifuScorer model from `{model_path}`")
-        
-        # Load MLP model
-        self.mlp = MLP(input_size=768)
-        s = torch.load(model_path, map_location=device)
-        self.mlp.load_state_dict(s)
-        self.mlp.to(device)
-        
-        # Load CLIP model
-        self.model2, self.preprocess = clip.load("ViT-L/14", device=device)
-        self.device = device
-        self.dtype = torch.float32
-        self.mlp.eval()
+
+        try:
+            import clip
+
+            if model_path is None:
+                model_path = "Eugeoter/waifu-scorer-v3/model.pth"
+                if self.verbose:
+                    print(f"model path not set, switch to default: `{model_path}`")
+
+            if not os.path.isfile(model_path):
+                split = model_path.split("/")
+                username, repo_id, model_name = split[-3], split[-2], split[-1]
+                model_path = hf_hub_download(f"{username}/{repo_id}", model_name, cache_dir=cache_dir)
+
+            print(f"Loading WaifuScorer model from `{model_path}`")
+
+            self.mlp = MLP(input_size=768)
+            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.model2, self.preprocess = clip.load("ViT-L/14", device=device)
+            self.device = device
+            self.dtype = torch.float32
+            self.mlp.eval()
+            self.available = True
+        except Exception as e:
+            print(f"Unable to initialize WaifuScorer: {e}")
+            self.available = False
 
     @torch.no_grad()
     def __call__(self, images):
+        if not self.available:
+            return [None] * (1 if not isinstance(images, list) else len(images))
+
         if isinstance(images, Image.Image):
             images = [images]
         n = len(images)
         if n == 1:
-            images = images*2  # batch norm requires at least 2 samples
-            
-        # Preprocess and encode images
+            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.model2.encode_image(image_batch)
-        
-        # Normalize features
+
         l2 = image_features.norm(2, dim=-1, keepdim=True)
         l2[l2 == 0] = 1
         im_emb_arr = (image_features / l2).to(device=self.device, dtype=self.dtype)
-        
-        # Get predictions
+
         predictions = self.mlp(im_emb_arr)
         scores = predictions.clamp(0, 10).cpu().numpy().reshape(-1).tolist()
-        
-        # Return only the requested number of scores
-        return scores[:n]
 
+        return scores[:n]
 
 def load_aesthetic_predictor_v2_5():
-    # This is a simplified version that just downloads the model
-    # The actual implementation would import and use aesthetic_predictor_v2_5
-    # We'll simulate the model with a dummy implementation
-    
-    class AestheticPredictorV2_5:
+    class AestheticPredictorV2_5_Impl: # Renamed class to avoid confusion
         def __init__(self):
             print("Loading Aesthetic Predictor V2.5...")
-            # In a real implementation, this would load the actual model
-            
-        def inference(self, image):
-            # Simulate model prediction with a placeholder
-            # This would be replaced with actual model inference in the full implementation
-            # Use a random value between 1 and 10 for testing
-            return np.random.uniform(1, 10)
-    
-    return AestheticPredictorV2_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: Image.Image) -> float:
+            # preprocess image
+            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()
+
+            # predict aesthetic score
+            with torch.inference_mode():
+                score = self.model(pixel_values).logits.squeeze().float().cpu().numpy()
 
+            return score
+
+    return AestheticPredictorV2_5_Impl() # Return an instance of the implementation class
 
 def load_anime_aesthetic_model():
     model_path = hf_hub_download(repo_id="skytnt/anime-aesthetic", filename="model.onnx")
     model = rt.InferenceSession(model_path, providers=['CPUExecutionProvider'])
     return model
 
-
 def predict_anime_aesthetic(img, model):
     img = np.array(img).astype(np.float32) / 255
     s = 768
@@ -140,224 +154,354 @@ def predict_anime_aesthetic(img, model):
     pred = model.run(None, {"img": img_input})[0].item()
     return pred
 
-
 class ImageEvaluationTool:
     def __init__(self):
         self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
         print(f"Using device: {self.device}")
-        
-        # Load all models
+
         print("Loading models... This may take some time.")
-        
-        # 1. Aesthetic Shadow
+
         print("Loading Aesthetic Shadow model...")
         self.aesthetic_shadow = pipeline("image-classification", model="NeoChen1024/aesthetic-shadow-v2-backup", device=self.device)
-        
-        try:
-            # 2. Waifu Scorer (requires CLIP)
-            print("Loading Waifu Scorer model...")
-            self.waifu_scorer = WaifuScorer(device=self.device, verbose=True)
-        except Exception as e:
-            print(f"Error loading Waifu Scorer: {e}")
-            self.waifu_scorer = None
-        
-        # 3. Aesthetic Predictor V2.5 (placeholder)
+
+        print("Loading Waifu Scorer model...")
+        self.waifu_scorer = WaifuScorer(device=self.device, verbose=True)
+
         print("Loading Aesthetic Predictor V2.5...")
         self.aesthetic_predictor_v2_5 = load_aesthetic_predictor_v2_5()
-        
-        # 4. Cafe Aesthetic models
-        print("Loading Cafe Aesthetic models...")
-        self.cafe_aesthetic = pipeline("image-classification", "cafeai/cafe_aesthetic")
-        self.cafe_style = pipeline("image-classification", "cafeai/cafe_style")
-        self.cafe_waifu = pipeline("image-classification", "cafeai/cafe_waifu")
-        
-        # 5. Anime Aesthetic
+
         print("Loading Anime Aesthetic model...")
         self.anime_aesthetic = load_anime_aesthetic_model()
-        
+
         print("All models loaded successfully!")
-        
-        # Create temp directory for storing processed images
+
         self.temp_dir = tempfile.mkdtemp()
-        
+
     def evaluate_image(self, image):
-        """Evaluate a single image with all models"""
         results = {}
-        
-        # Convert to PIL Image if not already
+
         if not isinstance(image, Image.Image):
             image = Image.fromarray(image)
-        
-        # 1. Aesthetic Shadow
+
         try:
             shadow_result = self.aesthetic_shadow(images=[image])[0]
             hq_score = [p for p in shadow_result if p['label'] == 'hq'][0]['score']
-            results['aesthetic_shadow'] = round(hq_score, 2)
+            # Scale aesthetic_shadow to 0-10 and clamp
+            aesthetic_shadow_score = np.clip(hq_score * 10.0, 0.0, 10.0)
+            results['aesthetic_shadow'] = aesthetic_shadow_score
         except Exception as e:
             print(f"Error in Aesthetic Shadow: {e}")
             results['aesthetic_shadow'] = None
-        
-        # 2. Waifu Scorer
-        if self.waifu_scorer:
-            try:
-                waifu_score = self.waifu_scorer([image])[0]
-                results['waifu_scorer'] = round(waifu_score, 2)
-            except Exception as e:
-                print(f"Error in Waifu Scorer: {e}")
-                results['waifu_scorer'] = None
-        else:
+
+        try:
+            waifu_score = self.waifu_scorer([image])[0]
+            # Clamp waifu_score
+            waifu_score_clamped = np.clip(waifu_score, 0.0, 10.0)
+            results['waifu_scorer'] = waifu_score_clamped
+        except Exception as e:
+            print(f"Error in Waifu Scorer: {e}")
             results['waifu_scorer'] = None
-        
-        # 3. Aesthetic Predictor V2.5
+
         try:
             v2_5_score = self.aesthetic_predictor_v2_5.inference(image)
-            results['aesthetic_predictor_v2_5'] = round(v2_5_score, 2)
+            # Clamp v2.5 score
+            v2_5_score_clamped = np.clip(v2_5_score, 0.0, 10.0)
+            results['aesthetic_predictor_v2_5'] = float(np.round(v2_5_score_clamped, 4)) # Keep 4 decimal places after clamping
         except Exception as e:
             print(f"Error in Aesthetic Predictor V2.5: {e}")
             results['aesthetic_predictor_v2_5'] = None
-        
-        # 4. Cafe Aesthetic
-        try:
-            cafe_aesthetic_result = self.cafe_aesthetic(image, top_k=2)
-            cafe_aesthetic_score = {d["label"]: round(d["score"], 2) for d in cafe_aesthetic_result}
-            results['cafe_aesthetic_good'] = cafe_aesthetic_score.get('good', 0)
-            results['cafe_aesthetic_bad'] = cafe_aesthetic_score.get('bad', 0)
-            
-            cafe_style_result = self.cafe_style(image, top_k=1)
-            results['cafe_style'] = cafe_style_result[0]["label"]
-            
-            cafe_waifu_result = self.cafe_waifu(image, top_k=1)
-            results['cafe_waifu'] = cafe_waifu_result[0]["label"]
-        except Exception as e:
-            print(f"Error in Cafe Aesthetic: {e}")
-            results['cafe_aesthetic_good'] = None
-            results['cafe_aesthetic_bad'] = None
-            results['cafe_style'] = None
-            results['cafe_waifu'] = None
-        
-        # 5. Anime Aesthetic
+
         try:
             img_array = np.array(image)
             anime_score = predict_anime_aesthetic(img_array, self.anime_aesthetic)
-            results['anime_aesthetic'] = round(anime_score, 2)
+            # Scale Anime Score to 0-10 and clamp
+            anime_score_scaled = np.clip(anime_score * 10.0, 0.0, 10.0)
+            results['anime_aesthetic'] = anime_score_scaled
         except Exception as e:
             print(f"Error in Anime Aesthetic: {e}")
             results['anime_aesthetic'] = None
-        
+
+        # Calculate Final Score (simple average of available scores)
+        valid_scores = [v for v in results.values() if v is not None]
+        if valid_scores:
+            final_score = np.mean(valid_scores)
+            results['final_score'] = np.clip(final_score, 0.0, 10.0) # Clamp final score too
+        else:
+            results['final_score'] = None
+
         return results
-    
-    def process_images(self, image_files):
-        """Process multiple image files and return results"""
+
+    def image_to_base64(self, image):
+        buffered = BytesIO()
+        image.save(buffered, format="JPEG")
+        return base64.b64encode(buffered.getvalue()).decode('utf-8')
+
+    def process_single_image(self, file_path):
+        try:
+            img = Image.open(file_path).convert("RGB")
+            eval_results = self.evaluate_image(img)
+            thumbnail = img.copy()
+            thumbnail.thumbnail((200, 200))
+            img_base64 = self.image_to_base64(thumbnail)
+            result = {
+                'file_name': os.path.basename(file_path),
+                'img_data': img_base64,
+                **eval_results
+            }
+            return result
+        except Exception as e:
+            print(f"Error processing {file_path}: {e}")
+            return None
+
+    def process_images_evaluation(self, image_files): # Renamed and now for evaluation only
         results = []
-        
+
         for i, file_path in enumerate(image_files):
             try:
-                # Open image
                 img = Image.open(file_path).convert("RGB")
-                
-                # Get image evaluation results
                 eval_results = self.evaluate_image(img)
-                
-                # Save a thumbnail for the results table
-                thumbnail_path = os.path.join(self.temp_dir, f"thumbnail_{i}.jpg")
-                img.thumbnail((200, 200))
-                img.save(thumbnail_path)
-                
-                # Add file info and thumbnail path to results
+
+                thumbnail = img.copy()
+                thumbnail.thumbnail((200, 200))
+
+                img_base64 = self.image_to_base64(thumbnail)
+
                 result = {
                     'file_name': os.path.basename(file_path),
-                    'thumbnail': thumbnail_path,
+                    'img_data': img_base64,
                     **eval_results
                 }
                 results.append(result)
-                
+
             except Exception as e:
                 print(f"Error processing {file_path}: {e}")
-        
+
+        return results
+
+    def sort_results(self, results, sort_by="Final Score"): # New function for sorting
+        def sort_key(res): # Define a sorting key function
+            sort_value = res.get(sort_by.lower().replace(" ", "_"), None) # Handle spaces and case
+            if sort_value is None: # Put N/A at the end
+                return -float('inf') if sort_by == "File Name" else float('inf') # File Name sort N/A at end alphabetically
+            return sort_value
+
+        results.sort(key=sort_key, reverse=sort_by != "File Name") # Sort results, reverse for score columns
         return results
-    
+
+    def generate_html_table(self, results):
+        html = """
+        <style>
+            .results-table {
+                width: 100%;
+                border-collapse: collapse;
+                margin: 20px 0;
+                font-family: Arial, sans-serif;
+                background-color: transparent;
+            }
+
+            .results-table th,
+            .results-table td {
+                color: #eee;
+                border: 1px solid #ddd;
+                padding: 8px;
+                text-align: center;
+                background-color: transparent;
+            }
+
+            .results-table th {
+                font-weight: bold;
+            }
+
+            .results-table tr:nth-child(even) {
+                background-color: transparent;
+            }
+
+            .results-table tr:hover {
+                background-color: rgba(255, 255, 255, 0.1);
+            }
+
+            .image-preview {
+                max-width: 150px;
+                max-height: 150px;
+                display: block;
+                margin: 0 auto;
+            }
+
+            .good-score {
+                color: #0f0;
+                font-weight: bold;
+            }
+            .bad-score {
+                color: #f00;
+                font-weight: bold;
+            }
+            .medium-score {
+                color: orange;
+                font-weight: bold;
+            }
+        </style>
+
+        <table class="results-table">
+            <thead>
+                <tr>
+                    <th>Image</th>
+                    <th>File Name</th>
+                    <th>Aesthetic Shadow</th>
+                    <th>Waifu Scorer</th>
+                    <th>Aesthetic V2.5</th>
+                    <th>Anime Score</th>
+                    <th>Final Score</th>
+                </tr>
+            </thead>
+            <tbody>
+        """
+
+        for result in results:
+            html += "<tr>"
+            html += f'<td><img src="data:image/jpeg;base64,{result["img_data"]}" class="image-preview"></td>'
+            html += f'<td>{result["file_name"]}</td>'
+
+            score = result["aesthetic_shadow"]
+            score_class = "good-score" if score and score >= 7 else "medium-score" if score and score >= 4 else "bad-score"
+            html += f'<td class="{score_class}">{score if score is not None else "N/A":.4f}</td>' # Format to 4 decimal places
+
+            score = result["waifu_scorer"]
+            score_class = "good-score" if score and score >= 7 else "medium-score" if score and score >= 5 else "bad-score"
+            html += f'<td class="{score_class}">{score if score is not None else "N/A":.4f}</td>' # Format to 4 decimal places
+
+            score = result["aesthetic_predictor_v2_5"]
+            score_class = "good-score" if score and score >= 7 else "medium-score" if score and score >= 5 else "bad-score"
+            html += f'<td class="{score_class}">{score if score is not None else "N/A":.4f}</td>' # Format to 4 decimal places
+
+            score = result["anime_aesthetic"]
+            score_class = "good-score" if score and score >= 7 else "medium-score" if score and score >= 5 else "bad-score"
+            html += f'<td class="{score_class}">{score if score is not None else "N/A":.4f}</td>' # Format to 4 decimal places
+
+            score = result["final_score"]
+            score_class = "good-score" if score and score >= 7 else "medium-score" if score and score >= 5 else "bad-score"
+            html += f'<td class="{score_class}">{score if score is not None else "N/A":.4f}</td>' # Format to 4 decimal places
+
+
+            html += "</tr>"
+
+        html += """
+            </tbody>
+        </table>
+        """
+
+        return html
+
     def cleanup(self):
-        """Clean up temporary files"""
         if os.path.exists(self.temp_dir):
             shutil.rmtree(self.temp_dir)
 
+# Global variable to store evaluation results
+global_results = None
 
-# Create the Gradio interface
 def create_interface():
+    global global_results # Use the global variable
+
     evaluator = ImageEvaluationTool()
-    
+    sort_options = ["Final Score", "File Name", "Aesthetic Shadow", "Waifu Scorer", "Aesthetic V2.5", "Anime Score"] # Sort options
+
     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:
-        
-        - **Aesthetic Shadow**: Evaluates high-quality vs low-quality images
+
+        - **Aesthetic Shadow**: Evaluates high-quality vs low-quality images (scaled to 0-10)
         - **Waifu Scorer**: Rates anime/illustration quality from 0-10
-        - **Aesthetic Predictor V2.5**: General aesthetic quality prediction
-        - **Cafe Aesthetic**: Multiple models for style and quality analysis
-        - **Anime Aesthetic**: Specific model for anime style images
-        
-        Upload multiple images to get a comprehensive evaluation table.
+        - **Aesthetic Predictor V2.5**: General aesthetic quality prediction (clamped to 0-10)
+        - **Anime Aesthetic**: Specific model for anime style images (scaled and clamped to 0-10)
+        - **Final Score**: Average of available scores (clamped to 0-10)
+
+        Upload multiple images to get a comprehensive evaluation table. Scores are clamped to the range 0.0000 - 10.0000.
         """)
-        
+
         with gr.Row():
             with gr.Column(scale=1):
                 input_images = gr.Files(label="Upload Images")
+                sort_dropdown = gr.Dropdown(sort_options, value="Final Score", label="Sort by") # Dropdown for sorting
                 process_btn = gr.Button("Evaluate Images", variant="primary")
                 clear_btn = gr.Button("Clear Results")
-            
+
             with gr.Column(scale=2):
-                output_gallery = gr.Gallery(label="Evaluated Images", columns=5, object_fit="contain")
-                output_table = gr.Dataframe(label="Evaluation Results")
-        
-        def process_images(files):
-            # Get file paths
+                progress_html = gr.HTML(label="Progress") # Keep progress_html if you want to show initial progress
+                output_html = gr.HTML(label="Evaluation Results")
+
+        def process_images_and_update(files): # Renamed and simplified
+            global global_results
             file_paths = [f.name for f in files]
-            
-            # Process images
-            results = evaluator.process_images(file_paths)
-            
-            # Prepare gallery and table
-            gallery_images = [{"image": r["thumbnail"], "label": f"{r['file_name']}"} for r in results]
-            
-            # Create DataFrame for the table
-            table_data = []
-            for r in results:
-                table_data.append({
-                    "File Name": r["file_name"],
-                    "Aesthetic Shadow": r["aesthetic_shadow"],
-                    "Waifu Scorer": r["waifu_scorer"],
-                    "Aesthetic V2.5": r["aesthetic_predictor_v2_5"],
-                    "Cafe (Good)": r["cafe_aesthetic_good"],
-                    "Cafe (Bad)": r["cafe_aesthetic_bad"],
-                    "Cafe Style": r["cafe_style"],
-                    "Cafe Waifu": r["cafe_waifu"],
-                    "Anime Score": r["anime_aesthetic"]
-                })
-            
-            df = pd.DataFrame(table_data)
-            return gallery_images, df
-        
+            total = len(file_paths)
+            progress_html_content = "" # Initialize progress content
+
+            if not file_paths: # Handle no files uploaded
+                global_results = []
+                return progress_html_content, evaluator.generate_html_table([]) # Empty table
+
+            progress_html_content = ""
+            for i, file_path in enumerate(file_paths):
+                percent = (i / total) * 100
+                progress_bar = f"""
+                <div>
+                    <p>Processing {os.path.basename(file_path)}</p>
+                    <progress value="{percent}" max="100"></progress>
+                    <p>{percent:.1f}% complete</p>
+                </div>
+                """
+                progress_html_content = progress_bar # Update progress content
+                yield progress_html_content, gr.update() # Yield progress update
+                # No need to process and sort here, just evaluate
+            global_results = evaluator.process_images_evaluation(file_paths) # Evaluate all images and store
+            sorted_results = evaluator.sort_results(global_results, sort_by="Final Score") # Initial sort by Final Score
+            html_table = evaluator.generate_html_table(sorted_results)
+            yield "<p>Processing complete</p>", html_table # Final progress and table
+
+        def update_table_sort(sort_by_column): # New function for sorting update
+            global global_results
+            if global_results is None:
+                return "No images evaluated yet." # Or handle case when no images are evaluated
+            sorted_results = evaluator.sort_results(global_results, sort_by=sort_by_column)
+            html_table = evaluator.generate_html_table(sorted_results)
+            return html_table
+
         def clear_results():
-            return None, None
-        
-        process_btn.click(process_images, inputs=[input_images], outputs=[output_gallery, output_table])
-        clear_btn.click(clear_results, inputs=[], outputs=[output_gallery, output_table])
-        
-        # Cleanup when closing
+            global global_results
+            global_results = None # Clear stored results
+            return gr.update(value=""), gr.update(value="")
+
+
+        process_btn.click(
+            process_images_and_update,
+            inputs=[input_images],
+            outputs=[progress_html, output_html]
+        )
+        sort_dropdown.change( # Only update table on sort change
+            update_table_sort,
+            inputs=[sort_dropdown],
+            outputs=[output_html] # Only update output_html
+        )
+        clear_btn.click(
+            clear_results,
+            inputs=[],
+            outputs=[progress_html, output_html]
+        )
+
         demo.load(lambda: None, inputs=None, outputs=None)
-        
+
         gr.Markdown("""
         ### Notes
         - The evaluation may take some time depending on the number and size of images
         - For best results, use high-quality images
-        - Scores are on different scales depending on the model
+        - Scores are color-coded: green for good (>=7), orange for medium (>=5), and red for poor scores (<5, or <4 for Aesthetic Shadow)
+        - Some models may fail for certain image types, shown as "N/A" in the results
+        - "Final Score" is a simple average of available model scores.
+        - Table is sortable by clicking the dropdown above the "Evaluate Images" button. Default sort is by "Final Score". Sorting happens instantly without re-evaluating images.
         """)
-    
+
     return demo
 
-# Launch the interface
 if __name__ == "__main__":
     demo = create_interface()
     demo.queue().launch()