Update app.py

app.py

@@ -19,91 +19,81 @@ except ImportError:
 
 # IQA-PyTorch imports
 try:
+    # This import needs to succeed for NIQE and MUSIQ
     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("ERROR: IQA-PyTorch library import failed. Some metrics (NIQE, MUSIQ-NR) will be unavailable. Check installation.")
+except ImportError as e:
+    print(f"ERROR: IQA-PyTorch library import failed: {e}. Some metrics (NIQE, MUSIQ-NR) will be unavailable. Check installation and dependencies (like kornia).")
+    IQA = None
+except Exception as e:
+    print(f"ERROR: An unexpected error occurred during IQA-PyTorch import: {e}")
     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 Normalization Parameters (Approximate typical ranges) ---
+# For "lower is better" metrics, score is (max_val - current_val) / (max_val - min_val)
+# For "higher is better" metrics, score is (current_val - min_val) / (max_val - min_val)
+# These are heuristics and can be adjusted.
+METRIC_RANGES = {
+    "brisque": {"min": 0, "max": 120, "lower_is_better": True}, # Typical BRISQUE range
+    "niqe": {"min": 0, "max": 12, "lower_is_better": True},    # Typical NIQE range
+    "musiq_nr": {"min": 10, "max": 90, "lower_is_better": False} # Example MUSIQ range
+}
+
 # --- 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:
         if img_tensor_chw_01.ndim == 3:
             img_tensor_bchw_01 = img_tensor_chw_01.unsqueeze(0)
         else:
             img_tensor_bchw_01 = img_tensor_chw_01
-
         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:
-        return f"Error"
-
+    except Exception: return "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
+        niqe_metric = IQA(libs='NIQE-PyTorch', device=DEVICE)
         score = niqe_metric(img_pil_rgb)
         return round(score.item(), 3)
-    except Exception as e:
-        return f"Error"
+    except Exception: return "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:
-        musiq_metric = IQA(libs='MUSIQ-L2N-lessons', device=DEVICE)
+        musiq_metric = IQA(libs='MUSIQ-L2N-lessons', device=DEVICE) # Example, could be other MUSIQ variants
         score = musiq_metric(img_pil_rgb)
         return round(score.item(), 3)
-    except Exception as e:
-        return f"Error"
-
+    except Exception: return "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:
-        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:
-             return f"FID needs at least 2 images per set. Found: Set1={len(set1_files)}, Set2={len(set2_files)}."
-
+        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: return f"FID needs at least 2 images per set. Found: Set1={len(set1_files)}, Set2={len(set2_files)}."
         fid_metric = piq.FID()
         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:
-             return "Feature extraction resulted in empty tensors."
-
+        if set1_features is None or set2_features is None: return "Could not extract features for one or both sets."
+        if set1_features.dim() == 0 or set2_features.dim() == 0 or set1_features.numel() == 0 or set2_features.numel() == 0: return "Feature extraction resulted in empty tensors."
         fid_value = fid_metric(set1_features, set2_features)
         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 ---
+# --- Helper & Final Score Calculation ---
 def pil_to_tensor_chw_01(img_pil_rgb):
     transform = T.Compose([T.ToTensor()])
     return transform(img_pil_rgb)
@@ -116,8 +106,39 @@ def create_thumbnail_base64(img_pil_rgb, size=THUMBNAIL_SIZE):
     img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
     return f"data:image/png;base64,{img_str}"
 
-# --- Main Processing Functions for Gradio ---
+def calculate_final_score(brisque_val, niqe_val, musiq_nr_val):
+    normalized_scores = []
+    
+    # BRISQUE
+    if isinstance(brisque_val, (float, int)):
+        cfg = METRIC_RANGES["brisque"]
+        val = max(cfg["min"], min(cfg["max"], brisque_val)) # Clip
+        norm_score = (cfg["max"] - val) / (cfg["max"] - cfg["min"]) if cfg["lower_is_better"] else (val - cfg["min"]) / (cfg["max"] - cfg["min"])
+        normalized_scores.append(norm_score)
+
+    # NIQE
+    if isinstance(niqe_val, (float, int)):
+        cfg = METRIC_RANGES["niqe"]
+        val = max(cfg["min"], min(cfg["max"], niqe_val)) # Clip
+        norm_score = (cfg["max"] - val) / (cfg["max"] - cfg["min"]) if cfg["lower_is_better"] else (val - cfg["min"]) / (cfg["max"] - cfg["min"])
+        normalized_scores.append(norm_score)
+
+    # MUSIQ-NR
+    if isinstance(musiq_nr_val, (float, int)):
+        cfg = METRIC_RANGES["musiq_nr"]
+        val = max(cfg["min"], min(cfg["max"], musiq_nr_val)) # Clip
+        norm_score = (cfg["max"] - val) / (cfg["max"] - cfg["min"]) if cfg["lower_is_better"] else (val - cfg["min"]) / (cfg["max"] - cfg["min"])
+        normalized_scores.append(norm_score)
+
+    if not normalized_scores:
+        return "N/A"
+    
+    # Average of normalized scores, then scale to 0-10
+    final_score_0_10 = (sum(normalized_scores) / len(normalized_scores)) * 10.0
+    return round(final_score_0_10, 4)
+
 
+# --- Main Processing Functions for Gradio ---
 def process_images_for_individual_scores(uploaded_file_list, progress=gr.Progress(track_tqdm=True)):
     if not uploaded_file_list:
         return pd.DataFrame(), "Please upload images first."
@@ -132,6 +153,7 @@ def process_images_for_individual_scores(uploaded_file_list, progress=gr.Progres
     results_data = []
 
     for i, file_obj in enumerate(uploaded_file_list):
+        base_filename = "Unknown File"
         try:
             file_path = file_obj.name
             base_filename = os.path.basename(file_path)
@@ -141,6 +163,9 @@ def process_images_for_individual_scores(uploaded_file_list, progress=gr.Progres
             brisque_val = get_brisque_score(img_tensor_chw_01)
             niqe_val = get_niqe_score(img_pil_rgb)
             musiq_nr_val = get_musiq_nr_score(img_pil_rgb)
+            
+            final_score = calculate_final_score(brisque_val, niqe_val, musiq_nr_val)
+            
             thumbnail_b64 = create_thumbnail_base64(img_pil_rgb)
             preview_html = f'<img src="{thumbnail_b64}" alt="{base_filename}">'
 
@@ -150,15 +175,15 @@ def process_images_for_individual_scores(uploaded_file_list, progress=gr.Progres
                 "BRISQUE (PIQ) (↓)": brisque_val,
                 "NIQE (IQA-PyTorch) (↓)": niqe_val,
                 "MUSIQ-NR (IQA-PyTorch) (↑)": musiq_nr_val,
+                "Final Score (0-10) (↑)": final_score,
             })
         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",
+                "BRISQUE (PIQ) (↓)": f"Load Err: {str(e)[:30]}",
                 "NIQE (IQA-PyTorch) (↓)": "N/A",
                 "MUSIQ-NR (IQA-PyTorch) (↑)": "N/A",
+                "Final Score (0-10) (↑)": "N/A",
             })
         progress((i + 1) / len(uploaded_file_list), desc=f"Processing {base_filename}")
 
@@ -166,38 +191,27 @@ def process_images_for_individual_scores(uploaded_file_list, progress=gr.Progres
     status_message += f"\nPer-image metrics calculated for {len(results_data)} images."
     return df_results, status_message
 
-
 def process_fid_for_two_sets(set1_file_list, set2_file_list, progress=gr.Progress(track_tqdm=True)):
     if not set1_file_list or not set2_file_list:
         return "Please upload files for both Set 1 and Set 2."
-
     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:
         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}."
-
+        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)}"
+    except Exception as e: fid_result_text = f"Error during FID preparation or calculation: {str(e)}"
     finally:
         if os.path.exists(set1_dir): shutil.rmtree(set1_dir)
         if os.path.exists(set2_dir): shutil.rmtree(set2_dir)
@@ -215,60 +229,42 @@ with gr.Blocks(theme=gr.themes.Soft(), css=css_custom) as demo:
     **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.
+    Final Score is a heuristic combination of available metrics (0-10, 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.")
-            
-            image_upload_input = gr.Files(
-                label=f"Upload Images (max {MAX_IMAGES_PER_BATCH}, .png, .jpg, .jpeg, .bmp, .webp)",
-                file_count="multiple",
-                type="filepath"
-            )
-            
+            gr.Markdown(f"Upload a batch of images (up to **{MAX_IMAGES_PER_BATCH}**) to get individual quality scores.")
+            image_upload_input = gr.Files(label=f"Upload Images (max {MAX_IMAGES_PER_BATCH}, .png, .jpg, .jpeg, .bmp, .webp)", file_count="multiple", type="filepath")
             evaluate_button_main = gr.Button("🖼️ Evaluate Uploaded Images", variant="primary")
-            
             gr.Markdown("---")
             status_output_main = gr.Textbox(label="📊 Evaluation Status", interactive=False, lines=2)
-
             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"],
+                headers=["Preview", "Filename", "BRISQUE (PIQ) (↓)", "NIQE (IQA-PyTorch) (↓)", "MUSIQ-NR (IQA-PyTorch) (↑)", "Final Score (0-10) (↑)"],
+                datatype=["html", "str", "number", "number", "number", "number"], # Added "number" for Final Score
                 interactive=False,
                 wrap=True,
-                row_count=(15, "paginate") # MODIFIED: Replaced max_rows and overflow_row_behaviour
+                row_count=(15, "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.
+            FID measures the similarity of two image distributions. **Lower FID scores are better**.
             """)
             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_set1_upload = gr.Files(label="Upload Images for Set 1", file_count="multiple", type="filepath")
+                fid_set2_upload = gr.Files(label="Upload Images for Set 2", 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)
 
-    evaluate_button_main.click(
-        fn=process_images_for_individual_scores,
-        inputs=[image_upload_input],
-        outputs=[results_table_output, status_output_main]
-    )
+    evaluate_button_main.click(fn=process_images_for_individual_scores, inputs=[image_upload_input], outputs=[results_table_output, status_output_main])
+    fid_calculate_button.click(fn=process_fid_for_two_sets, inputs=[fid_set1_upload, fid_set2_upload], outputs=[fid_result_output])
 
-    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:
+# --- For Hugging Face Spaces ---
+# Ensure 'requirements.txt' includes:
 """
 gradio
 torch
@@ -276,18 +272,14 @@ torchvision
 Pillow
 numpy
 piq>=0.8.0
-iqa-pytorch==0.2.1 # PINNED VERSION
+iqa-pytorch==0.1
 timm
 scikit-image
 pandas
+kornia 
 """
 
 if __name__ == "__main__":
-    if piq is None:
-        print("\n\nWARNING: PIQ library is missing.")
-        print("Please install it: pip install piq\n\n")
-    if IQA is None:
-        print("\n\nERROR: IQA-PyTorch library import failed or it's missing.")
-        print("Please ensure it's installed correctly (e.g., pip install iqa-pytorch==0.2.1) and check for import errors during startup.\n\n")
-    
+    if piq is None: print("\nWARNING: PIQ library is missing. pip install piq\n")
+    if IQA is None: print("\nERROR: IQA-PyTorch library import failed. pip install iqa-pytorch==0.1 kornia\n")
     demo.launch(debug=True)