Update app.py

app.py

@@ -430,33 +430,164 @@ def create_interface():
                 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
+        def process_images_and_update(files):
             global global_results
             file_paths = [f.name for f in files]
-            total = len(file_paths)
-            progress_html_content = "" # Initialize progress content
-
-            if not file_paths: # Handle no files uploaded
+            total_files = len(file_paths)
+            results = []
+            
+            if not file_paths:
                 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>
+                yield "<p>No files uploaded.</p>", gr.update()
+                return
+
+            # Helper function to generate a styled progress bar HTML snippet.
+            def generate_progress_bar(percentage):
+                return f"""
+                <div style="background-color: #ddd; border-radius: 5px; width: 100%; margin: 10px 0;">
+                <div style="width: {percentage:.1f}%; background-color: #4CAF50; text-align: center; padding: 5px 0; border-radius: 5px;">
+                    {percentage:.1f}%
+                </div>
                 </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
+            
+            total_models = 4  # Total number of model steps per image.
+            
+            for i, file_path in enumerate(file_paths):
+                file_name = os.path.basename(file_path)
+                try:
+                    img = Image.open(file_path).convert("RGB")
+                except Exception as e:
+                    yield f"<p>Error opening {file_name}: {e}</p>", gr.update()
+                    continue
+                
+                # Update overall progress before starting this image.
+                overall_percent = (i / total_files) * 100
+                overall_bar = generate_progress_bar(overall_percent)
+                progress_html = f"""
+                <p>{file_name}: Starting evaluation...</p>
+                <p>Overall Progress:</p>
+                {overall_bar}
+                """
+                yield progress_html, gr.update()
+                
+                # === Model Step 1: Aesthetic Shadow ===
+                model_index = 0
+                sub_percent = ((model_index + 1) / total_models) * 100
+                sub_bar = generate_progress_bar(sub_percent)
+                progress_html = f"""
+                <p>{file_name}: Evaluating <strong>Aesthetic Shadow</strong>...</p>
+                <p>Current Image Progress:</p>
+                {sub_bar}
+                <p>Overall Progress:</p>
+                {overall_bar}
+                """
+                yield progress_html, gr.update()
+                try:
+                    shadow_result = evaluator.aesthetic_shadow(images=[img])[0]
+                    hq_score = [p for p in shadow_result if p['label'] == 'hq'][0]['score']
+                    aesthetic_shadow_score = np.clip(hq_score * 10.0, 0.0, 10.0)
+                except Exception as e:
+                    print(f"Error in Aesthetic Shadow for {file_name}: {e}")
+                    aesthetic_shadow_score = None
+                yield f"<p>{file_name}: <strong>Aesthetic Shadow</strong> evaluation complete.</p>", gr.update()
+                
+                # === Model Step 2: Waifu Scorer ===
+                model_index = 1
+                sub_percent = ((model_index + 1) / total_models) * 100
+                sub_bar = generate_progress_bar(sub_percent)
+                progress_html = f"""
+                <p>{file_name}: Evaluating <strong>Waifu Scorer</strong>...</p>
+                <p>Current Image Progress:</p>
+                {sub_bar}
+                <p>Overall Progress:</p>
+                {overall_bar}
+                """
+                yield progress_html, gr.update()
+                try:
+                    waifu_score = evaluator.waifu_scorer([img])[0]
+                    waifu_score = np.clip(waifu_score, 0.0, 10.0)
+                except Exception as e:
+                    print(f"Error in Waifu Scorer for {file_name}: {e}")
+                    waifu_score = None
+                yield f"<p>{file_name}: <strong>Waifu Scorer</strong> evaluation complete.</p>", gr.update()
+                
+                # === Model Step 3: Aesthetic Predictor V2.5 ===
+                model_index = 2
+                sub_percent = ((model_index + 1) / total_models) * 100
+                sub_bar = generate_progress_bar(sub_percent)
+                progress_html = f"""
+                <p>{file_name}: Evaluating <strong>Aesthetic Predictor V2.5</strong>...</p>
+                <p>Current Image Progress:</p>
+                {sub_bar}
+                <p>Overall Progress:</p>
+                {overall_bar}
+                """
+                yield progress_html, gr.update()
+                try:
+                    v2_5_score = evaluator.aesthetic_predictor_v2_5.inference(img)
+                    v2_5_score = float(np.round(np.clip(v2_5_score, 0.0, 10.0), 4))
+                except Exception as e:
+                    print(f"Error in Aesthetic Predictor V2.5 for {file_name}: {e}")
+                    v2_5_score = None
+                yield f"<p>{file_name}: <strong>Aesthetic Predictor V2.5</strong> evaluation complete.</p>", gr.update()
+                
+                # === Model Step 4: Anime Aesthetic ===
+                model_index = 3
+                sub_percent = ((model_index + 1) / total_models) * 100
+                sub_bar = generate_progress_bar(sub_percent)
+                progress_html = f"""
+                <p>{file_name}: Evaluating <strong>Anime Aesthetic</strong>...</p>
+                <p>Current Image Progress:</p>
+                {sub_bar}
+                <p>Overall Progress:</p>
+                {overall_bar}
+                """
+                yield progress_html, gr.update()
+                try:
+                    img_array = np.array(img)
+                    anime_score = predict_anime_aesthetic(img_array, evaluator.anime_aesthetic)
+                    anime_score = np.clip(anime_score * 10.0, 0.0, 10.0)
+                except Exception as e:
+                    print(f"Error in Anime Aesthetic for {file_name}: {e}")
+                    anime_score = None
+                yield f"<p>{file_name}: <strong>Anime Aesthetic</strong> evaluation complete.</p>", gr.update()
+                
+                # === Final Score Calculation and Results Collection ===
+                valid_scores = [v for v in [aesthetic_shadow_score, waifu_score, v2_5_score, anime_score] if v is not None]
+                final_score = np.clip(np.mean(valid_scores), 0.0, 10.0) if valid_scores else None
+                
+                # Create a thumbnail and store the evaluation results.
+                thumbnail = img.copy()
+                thumbnail.thumbnail((200, 200))
+                img_base64 = evaluator.image_to_base64(thumbnail)
+                result = {
+                    'file_name': file_name,
+                    'img_data': img_base64,
+                    'aesthetic_shadow': aesthetic_shadow_score,
+                    'waifu_scorer': waifu_score,
+                    'aesthetic_predictor_v2_5': v2_5_score,
+                    'anime_aesthetic': anime_score,
+                    'final_score': final_score
+                }
+                results.append(result)
+                
+                # Update overall progress for the processed file.
+                overall_percent = ((i + 1) / total_files) * 100
+                overall_bar = generate_progress_bar(overall_percent)
+                progress_html = f"""
+                <p>{file_name}: Evaluation complete. ({i + 1}/{total_files} images processed.)</p>
+                <p>Overall Progress:</p>
+                {overall_bar}
+                """
+                # Sort the results by Final Score and update the table.
+                sorted_results = evaluator.sort_results(results.copy(), sort_by="Final Score")
+                html_table = evaluator.generate_html_table(sorted_results)
+                yield progress_html, html_table
+            
+            # Sort final results by Final Score.
+            global_results = evaluator.sort_results(results, sort_by="Final Score")
+            yield "<p>All images processed.</p>", evaluator.generate_html_table(global_results)
 
         def update_table_sort(sort_by_column): # New function for sorting update
             global global_results