Update app.py

app.py

@@ -4,7 +4,6 @@ import torch
 import torch.nn as nn
 from PIL import Image
 import torchvision.transforms as transforms
-from pathlib import Path
 import numpy as np
 from huggingface_hub import hf_hub_download
 
@@ -49,7 +48,6 @@ class CompressionArtifactPredictor:
         checkpoint = torch.load(model_path, map_location=self.device, weights_only=True)
         self.model.load_state_dict(checkpoint['model_state_dict'])
 
-        # Define preprocessing
         self.preprocess = transforms.Compose([
             transforms.ToTensor(),
         ])
@@ -63,21 +61,17 @@ class CompressionArtifactPredictor:
         }
 
     def predict(self, image: Image.Image) -> dict:
-        """Predict compression quality levels for all formats with explanations."""
-        # Preprocess
+        """Predict compression quality levels for all formats."""
         img_tensor = self.preprocess(image).unsqueeze(0).to(self.device)
 
-        # Inference
         with torch.no_grad():
             with torch.cuda.amp.autocast(dtype=torch.bfloat16):
                 predictions = self.model(img_tensor).squeeze(0).cpu().float().numpy()
 
-        # Format results
         results = {}
         for i, fmt in enumerate(self.compression_formats):
-            quality_score = float(predictions[i] * 100)  # Convert to 0-100 range
+            quality_score = float(predictions[i] * 100)
             
-            # Determine quality category
             if quality_score >= 90:
                 category = "Excellent (Minimal artifacts)"
                 color = "🟢"
@@ -109,17 +103,12 @@ def create_ui():
         if image is None:
             return None, "Please upload an image."
         
-        # Convert numpy array to PIL Image if needed
         if isinstance(image, np.ndarray):
             image = Image.fromarray(image)
         
-        # Ensure RGB
         image = image.convert('RGB')
-        
-        # Get predictions
         results = predictor.predict(image)
         
-        # Format output as a nice dictionary for Gradio
         formatted_results = {}
         for fmt, data in results.items():
             formatted_results[f"{data['indicator']} {fmt}"] = {
@@ -128,21 +117,18 @@ def create_ui():
                 "Model Accuracy": f"{data['accuracy']}%"
             }
         
-        # Calculate overall score
         avg_quality = np.mean([r['quality_score'] for r in results.values()])
         if avg_quality >= 85:
-            overall_status = "✅ **High Quality Image** - Minimal compression artifacts detected across all formats."
+            overall_status = "✅ **High Quality Image** - Minimal compression artifacts detected."
         elif avg_quality >= 65:
-            overall_status = "⚠️ **Moderate Quality** - Some compression artifacts present, but image remains usable."
+            overall_status = "⚠️ **Moderate Quality** - Some compression artifacts present, but usable."
         else:
             overall_status = "❌ **Low Quality Image** - Significant compression artifacts detected."
         
-        # Add overall summary
         summary = f"### Overall Assessment\n{overall_status}\n\n**Average Quality Score: {avg_quality:.1f}/100**"
         
         return formatted_results, summary
 
-    # Create Gradio interface
     with gr.Blocks(
         title="AAL-Plus Image Quality Assessment",
         theme=gr.themes.Soft()
@@ -150,19 +136,13 @@ def create_ui():
         gr.Markdown(
             """
             # 🎯 AAL-Plus Image Quality Assessment
-            ### Detect compression artifacts across multiple formats (JPEG, WebP, AVIF, JXL)
+            ### Detect compression artifacts across multiple image formats (JPEG, WebP, AVIF, JXL)
             
-            This lightweight model (~2M parameters) predicts the quality level of your image and identifies 
-            compression artifacts with **97.1% overall accuracy**.
+            This lightweight model (~2M parameters, 8MB) predicts quality levels with **97.1% overall accuracy**.
             
             **How to interpret results:**
-            - **Quality Score**: 0-100 scale (higher = better quality, fewer artifacts)
-            - **Assessment**: Text description of artifact level
-            - **Color Indicators**: 
-                - 🟢 Green = Excellent (90-100)
-                - 🟡 Yellow = Good (70-90)
-                - 🟠 Orange = Fair (50-70)
-                - 🔴 Red = Poor (0-50)
+            - **Quality Score**: 0-100 scale (higher = better quality)
+            - **Score Categories**: 🟢 90-100 | 🟡 70-90 | 🟠 50-70 | 🔴 0-50
             """
         )
         
@@ -173,7 +153,7 @@ def create_ui():
                     type="pil",
                     height=400
                 )
-                analyze_button = gr.Button("🔍 Analyze Image Quality", variant="primary")
+                analyze_button = gr.Button("🔍 Analyze Image Quality", variant="primary", size="lg")
             
             with gr.Column():
                 results_output = gr.Label(
@@ -184,21 +164,6 @@ def create_ui():
                     label="Overall Assessment"
                 )
         
-        # Examples
-        gr.Examples(
-            examples=[
-                ["examples/example1.jpg"],
-                ["examples/example2.webp"],
-                ["examples/example3.avif"],
-                ["examples/example4.jxl"],
-            ],
-            inputs=image_input,
-            outputs=[results_output, summary_output],
-            fn=analyze_image,
-            cache_examples=False,
-            label="Try Example Images"
-        )
-        
         gr.Markdown(
             """
             ---
@@ -210,13 +175,10 @@ def create_ui():
             | AVIF   | 97.1%              | 0-100         |
             | JXL    | 94.8%              | 0-100         |
             
-            *Accuracy measured as predictions within ±5% range of actual quality values*
-            
-            **🌍 Environmental Impact**: Model training required ~12 GPU hours on RTX 5090. Model size: 8MB.
+            *Accuracy measured as predictions within ±5% of actual quality values*
             """
         )
         
-        # Wire up the interface
         analyze_button.click(
             fn=analyze_image,
             inputs=image_input,