Create app.py

app.py

@@ -0,0 +1,332 @@
+import gradio as gr
+import pandas as pd
+import torch
+import torchvision.transforms as T
+from PIL import Image
+import numpy as np
+import io
+import base64
+import os
+import shutil
+import tempfile
+
+# PIQ imports
+try:
+    import piq
+except ImportError:
+    print("Warning: PIQ library not found. Some metrics (BRISQUE, FID) will be unavailable.")
+    piq = None
+
+# IQA-PyTorch imports
+try:
+    from iqa_pytorch import IQA
+    # Available models in IQA-PyTorch (examples for NR):
+    # "MUSIQ-L2N-lessons", "MUSIQ-Koniq-NSR", "MUSIQ-SpAq-NSR"
+    # "BRISQUE-PyTorch", "NIQE-PyTorch"
+    # "NIMA-VGG16-estimate", "NIMA-MobileNet-estimate" (Aesthetic)
+except ImportError:
+    print("Warning: IQA-PyTorch library not found. Some metrics (NIQE, MUSIQ-NR) will be unavailable.")
+    IQA = None
+
+# --- Configuration ---
+DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+MAX_IMAGES_PER_BATCH = 100
+THUMBNAIL_SIZE = (64, 64) # (width, height) for preview
+
+# --- Metric Functions ---
+
+def get_brisque_score(img_tensor_chw_01):
+    """Calculates BRISQUE score using PIQ. Expects a (C, H, W) tensor, range [0, 1]."""
+    if piq is None: return "N/A (PIQ missing)"
+    try:
+        # Ensure tensor is (B, C, H, W) for piq.brisque
+        if img_tensor_chw_01.ndim == 3:
+            img_tensor_bchw_01 = img_tensor_chw_01.unsqueeze(0)
+        else: # Already has batch dim or incorrect dims
+            img_tensor_bchw_01 = img_tensor_chw_01
+
+        # Ensure 3 channels if it's grayscale by repeating
+        if img_tensor_bchw_01.shape[1] == 1:
+            img_tensor_bchw_01 = img_tensor_bchw_01.repeat(1, 3, 1, 1)
+
+        brisque_loss = piq.brisque(img_tensor_bchw_01.to(DEVICE), data_range=1.)
+        return round(brisque_loss.item(), 3)
+    except Exception as e:
+        # print(f"BRISQUE Error: {e} for tensor shape {img_tensor_chw_01.shape}")
+        return f"Error"
+
+
+def get_niqe_score(img_pil_rgb):
+    """Calculates NIQE score using IQA-PyTorch. Expects a PIL RGB image."""
+    if IQA is None: return "N/A (IQA missing)"
+    try:
+        niqe_metric = IQA(libs='NIQE-PyTorch', device=DEVICE) # NIQE is No-Reference
+        score = niqe_metric(img_pil_rgb)
+        return round(score.item(), 3)
+    except Exception as e:
+        # print(f"NIQE Error: {e}")
+        return f"Error"
+
+def get_musiq_nr_score(img_pil_rgb):
+    """Calculates No-Reference MUSIQ score using IQA-PyTorch. Expects a PIL RGB image."""
+    if IQA is None: return "N/A (IQA missing)"
+    try:
+        # Using MUSIQ-L2N-lessons as an example NR model from IQA-PyTorch
+        # Other options: "MUSIQ-Koniq-NSR", "MUSIQ-SpAq-NSR"
+        musiq_metric = IQA(libs='MUSIQ-L2N-lessons', device=DEVICE)
+        score = musiq_metric(img_pil_rgb)
+        return round(score.item(), 3)
+    except Exception as e:
+        # print(f"MUSIQ-NR Error: {e}")
+        return f"Error"
+
+
+def get_fid_score_piq_folders(path_to_set1_folder, path_to_set2_folder):
+    """Calculates FID between two folders of images using PIQ."""
+    if piq is None: return "N/A (PIQ missing)"
+    try:
+        # List image files in folders
+        set1_files = [os.path.join(path_to_set1_folder, f) for f in os.listdir(path_to_set1_folder)
+                      if f.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp', '.webp'))]
+        set2_files = [os.path.join(path_to_set2_folder, f) for f in os.listdir(path_to_set2_folder)
+                      if f.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp', '.webp'))]
+
+        if not set1_files or not set2_files:
+            return "One or both sets have no valid image files."
+        if len(set1_files) < 2 or len(set2_files) < 2: # FID usually needs more, but PIQ might handle small N. Min 2 to compute stats.
+             return f"FID needs at least 2 images per set. Found: Set1={len(set1_files)}, Set2={len(set2_files)}."
+
+        fid_metric = piq.FID()
+        # compute_feats expects a list of image paths
+        set1_features = fid_metric.compute_feats(set1_files, device=DEVICE)
+        set2_features = fid_metric.compute_feats(set2_files, device=DEVICE)
+
+        if set1_features is None or set2_features is None:
+             return "Could not extract features for one or both sets (check image validity and count)."
+        if set1_features.dim() == 0 or set2_features.dim() == 0 or set1_features.numel() == 0 or set2_features.numel() == 0: # Handle empty tensors
+             return "Feature extraction resulted in empty tensors."
+
+
+        fid_value = fid_metric(set1_features, set2_features) # Pass tensors directly
+        return round(fid_value.item(), 3)
+    except Exception as e:
+        print(f"FID calculation error: {e}")
+        return f"FID Error: {str(e)[:100]}"
+
+# --- Helper Functions ---
+def pil_to_tensor_chw_01(img_pil_rgb):
+    """Converts PIL RGB image to PyTorch CHW tensor [0,1]."""
+    transform = T.Compose([T.ToTensor()]) # Converts PIL [0,255] to Tensor [0,1] C,H,W
+    return transform(img_pil_rgb)
+
+def create_thumbnail_base64(img_pil_rgb, size=THUMBNAIL_SIZE):
+    """Creates a base64 encoded PNG thumbnail string from a PIL image."""
+    img_copy = img_pil_rgb.copy()
+    img_copy.thumbnail(size)
+    buffered = io.BytesIO()
+    img_copy.save(buffered, format="PNG")
+    img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
+    return f"data:image/png;base64,{img_str}"
+
+# --- Main Processing Functions for Gradio ---
+
+def process_images_for_individual_scores(uploaded_file_list, progress=gr.Progress(track_tqdm=True)):
+    """Processes uploaded images for individual quality scores and displays them."""
+    if not uploaded_file_list:
+        return pd.DataFrame(), "Please upload images first.", "IS: N/A (Not Implemented)", "FID: N/A (Use FID Tab)"
+
+    # Limit number of images
+    if len(uploaded_file_list) > MAX_IMAGES_PER_BATCH:
+         status_message = f"Too many images ({len(uploaded_file_list)}). Processing first {MAX_IMAGES_PER_BATCH} images."
+         uploaded_file_list = uploaded_file_list[:MAX_IMAGES_PER_BATCH]
+    else:
+        status_message = f"Processing {len(uploaded_file_list)} images..."
+    
+    progress(0, desc=status_message)
+
+    results_data = []
+    # Temporary directory for this batch if needed by some metric that takes a folder path
+    # batch_temp_dir = tempfile.mkdtemp(prefix="eval_batch_")
+
+    for i, file_obj in enumerate(uploaded_file_list):
+        try:
+            # file_obj for gr.Files is a tempfile._TemporaryFileWrapper object
+            file_path = file_obj.name
+            base_filename = os.path.basename(file_path)
+
+            img_pil_rgb = Image.open(file_path).convert("RGB")
+
+            # 1. For PIQ BRISQUE (needs tensor)
+            img_tensor_chw_01 = pil_to_tensor_chw_01(img_pil_rgb)
+            brisque_val = get_brisque_score(img_tensor_chw_01)
+
+            # 2. For IQA-PyTorch NIQE & MUSIQ (needs PIL image)
+            niqe_val = get_niqe_score(img_pil_rgb)
+            musiq_nr_val = get_musiq_nr_score(img_pil_rgb)
+
+            # 3. Thumbnail for display
+            thumbnail_b64 = create_thumbnail_base64(img_pil_rgb)
+            preview_html = f'<img src="{thumbnail_b64}" alt="{base_filename}">'
+
+            results_data.append({
+                "Preview": preview_html,
+                "Filename": base_filename,
+                "BRISQUE (PIQ) (↓)": brisque_val,
+                "NIQE (IQA-PyTorch) (↓)": niqe_val,
+                "MUSIQ-NR (IQA-PyTorch) (↑)": musiq_nr_val,
+            })
+        except Exception as e:
+            try: base_filename = os.path.basename(file_obj.name if hasattr(file_obj, 'name') else str(file_obj))
+            except: base_filename = "Unknown File"
+            results_data.append({
+                "Preview": "Error processing", "Filename": base_filename,
+                "BRISQUE (PIQ) (↓)": f"Load Err",
+                "NIQE (IQA-PyTorch) (↓)": "N/A",
+                "MUSIQ-NR (IQA-PyTorch) (↑)": "N/A",
+            })
+        progress((i + 1) / len(uploaded_file_list), desc=f"Processing {base_filename}")
+
+    df_results = pd.DataFrame(results_data)
+    status_message += f"\nPer-image metrics calculated for {len(results_data)} images."
+
+    # Batch metrics info (IS not implemented, FID separate)
+    is_text = "IS (PIQ): Not implemented in this version."
+    fid_text_batch_info = "FID (PIQ): Use the 'Calculate FID (Set vs Set)' tab for FID scores."
+
+    # Cleanup temp dir if created
+    # if os.path.exists(batch_temp_dir): shutil.rmtree(batch_temp_dir)
+
+    return df_results, status_message, is_text, fid_text_batch_info
+
+
+def process_fid_for_two_sets(set1_file_list, set2_file_list, progress=gr.Progress(track_tqdm=True)):
+    """Handles FID calculation between two sets of uploaded images."""
+    if not set1_file_list or not set2_file_list:
+        return "Please upload files for both Set 1 and Set 2."
+
+    # Create temporary directories for Set 1 and Set 2
+    # Suffix helps identify user folders if many users hit it, though Gradio handles sessions.
+    # Prefix is better for mkdtemp.
+    set1_dir = tempfile.mkdtemp(prefix="fid_set1_")
+    set2_dir = tempfile.mkdtemp(prefix="fid_set2_")
+    
+    fid_result_text = "Starting FID calculation..."
+    progress(0.1, desc="Preparing image sets for FID...")
+
+    try:
+        # Copy uploaded files to these temporary directories
+        for i, file_obj in enumerate(set1_file_list):
+            shutil.copy(file_obj.name, os.path.join(set1_dir, os.path.basename(file_obj.name)))
+            progress(0.1 + 0.2 * (i / len(set1_file_list)), desc=f"Copying Set 1: {os.path.basename(file_obj.name)}")
+
+        for i, file_obj in enumerate(set2_file_list):
+            shutil.copy(file_obj.name, os.path.join(set2_dir, os.path.basename(file_obj.name)))
+            progress(0.3 + 0.2 * (i / len(set2_file_list)), desc=f"Copying Set 2: {os.path.basename(file_obj.name)}")
+        
+        num_set1 = len(os.listdir(set1_dir))
+        num_set2 = len(os.listdir(set2_dir))
+
+        if num_set1 == 0 or num_set2 == 0:
+             return f"FID Error: One or both sets are empty after copying. Set 1: {num_set1}, Set 2: {num_set2}."
+
+        progress(0.5, desc=f"Calculating FID between Set 1 ({num_set1} images) and Set 2 ({num_set2} images)...")
+        fid_score = get_fid_score_piq_folders(set1_dir, set2_dir)
+        progress(1, desc="FID calculation complete.")
+        fid_result_text = f"FID (PIQ) between Set 1 ({num_set1} images) and Set 2 ({num_set2} images): {fid_score}"
+
+    except Exception as e:
+        fid_result_text = f"Error during FID preparation or calculation: {str(e)}"
+    finally:
+        # Cleanup temporary directories
+        if os.path.exists(set1_dir): shutil.rmtree(set1_dir)
+        if os.path.exists(set2_dir): shutil.rmtree(set2_dir)
+
+    return fid_result_text
+
+# --- Gradio UI Definition ---
+css_custom = """
+table {font-size: 0.8em !important; width: 100% !important;} 
+th, td {padding: 4px !important; text-align: left !important;}
+img {max-width: 64px !important; max-height: 64px !important; object-fit: contain;}
+"""
+with gr.Blocks(theme=gr.themes.Soft(), css=css_custom) as demo:
+    gr.Markdown(f"""
+    # Image Generation Model Evaluation Tool
+    **Objective:** Automated evaluation and comparison of image quality from different model versions.
+    Utilizes `PIQ` and `IQA-PyTorch` libraries. Runs on **{DEVICE}**.
+    (↓) means lower is better, (↑) means higher is better.
+    """)
+
+    with gr.Tabs():
+        with gr.TabItem("Per-Image Quality Evaluation"):
+            gr.Markdown(f"Upload a batch of images (up to **{MAX_IMAGES_PER_BATCH}**) to get individual quality scores. Images are processed in the browser's session.")
+            
+            # Using gr.Files which allows multiple uploads and returns a list of TemporaryFileWrapper objects
+            image_upload_input = gr.Files(
+                label=f"Upload Images (max {MAX_IMAGES_PER_BATCH}, .png, .jpg, .jpeg, .bmp, .webp)",
+                file_count="multiple",
+                type="filepath" # Provides path to temp file
+            )
+            
+            evaluate_button_main = gr.Button("🖼️ Evaluate Uploaded Images", variant="primary")
+            
+            gr.Markdown("---")
+            status_output_main = gr.Textbox(label="📊 Evaluation Status", interactive=False, lines=2)
+            # batch_is_output = gr.Textbox(label="Overall Inception Score (IS) for Batch", interactive=False, lines=1) # IS deferred
+            # batch_fid_output_info = gr.Textbox(label="Overall Fréchet Inception Distance (FID) for Batch", interactive=False, lines=1) # FID separate
+
+            gr.Markdown("### 🖼️ Per-Image Evaluation Results")
+            gr.Markdown("Click column headers to sort. Previews are thumbnails.")
+            results_table_output = gr.DataFrame(
+                headers=["Preview", "Filename", "BRISQUE (PIQ) (↓)", "NIQE (IQA-PyTorch) (↓)", "MUSIQ-NR (IQA-PyTorch) (↑)"],
+                datatype=["html", "str", "number", "number", "number"],
+                interactive=False, wrap=True, max_rows=15, overflow_row_behaviour="paginate"
+            )
+
+        with gr.TabItem("↔️ Calculate FID (Set vs. Set)"):
+            gr.Markdown("""
+            Calculate Fréchet Inception Distance (FID) between two sets of images. 
+            FID measures the similarity of two image distributions (e.g., generated vs. real, or version A vs. version B). 
+            **Lower FID scores are better**, indicating more similarity.
+            """)
+            with gr.Row():
+                fid_set1_upload = gr.Files(label="Upload Images for Set 1 (.png, .jpg, .jpeg, .bmp, .webp)", file_count="multiple", type="filepath")
+                fid_set2_upload = gr.Files(label="Upload Images for Set 2 (.png, .jpg, .jpeg, .bmp, .webp)", file_count="multiple", type="filepath")
+            
+            fid_calculate_button = gr.Button("🔗 Calculate FID between Set 1 and Set 2", variant="primary")
+            fid_result_output = gr.Textbox(label="📈 FID Result", interactive=False, lines=2)
+
+    # Wire components
+    evaluate_button_main.click(
+        fn=process_images_for_individual_scores,
+        inputs=[image_upload_input],
+        outputs=[results_table_output, status_output_main] #, batch_is_output, batch_fid_output_info]
+    )
+
+    fid_calculate_button.click(
+        fn=process_fid_for_two_sets,
+        inputs=[fid_set1_upload, fid_set2_upload],
+        outputs=[fid_result_output]
+    )
+
+# --- For Hugging Face Spaces: requirements.txt ---
+# Ensure this content is in your 'requirements.txt' file in the HF Space:
+"""
+gradio
+torch
+torchvision
+Pillow
+numpy
+piq>=0.8.0 # Specify version if known good, or just piq
+iqa-pytorch>=0.2.1 # Specify version if known good
+timm # A dependency for some iqa-pytorch models like MUSIQ
+scikit-image # Often a transitive dependency, good to include
+pandas
+"""
+
+if __name__ == "__main__":
+    if piq is None or IQA is None:
+        print("\n\nWARNING: One or more core libraries (PIQ, IQA-PyTorch) are missing.")
+        print("Please install them by creating a 'requirements.txt' file with the content above and running: pip install -r requirements.txt\n\n")
+    
+    demo.launch(debug=True) # Set debug=False for production