Update app.py

app.py

@@ -1,4 +1,5 @@
 # app.py
+import warnings
 import gradio as gr
 import torch
 import torch.nn as nn
@@ -7,6 +8,9 @@ import torchvision.transforms as transforms
 import numpy as np
 from huggingface_hub import hf_hub_download
 
+# Suppress HF spaces warning (internal library, not our code)
+warnings.filterwarnings("ignore", category=FutureWarning, module="spaces")
+
 # ==================== MODEL DEFINITION ====================
 class LightweightCompressionNet(nn.Module):
     def __init__(self):
@@ -31,7 +35,6 @@ class LightweightCompressionNet(nn.Module):
         features = features.view(features.size(0), -1)
         return self.head(features)
 
-
 # ==================== INFERENCE CLASS ====================
 class CompressionArtifactPredictor:
     def __init__(self, device: str = "cuda"):
@@ -48,8 +51,11 @@ class CompressionArtifactPredictor:
         checkpoint = torch.load(model_path, map_location=self.device, weights_only=True)
         self.model.load_state_dict(checkpoint['model_state_dict'])
 
+        # FIXED: Add padding->center crop to handle arbitrary sizes
         self.preprocess = transforms.Compose([
             transforms.ToTensor(),
+            transforms.Pad(512, padding_mode='edge'),  # Pad smaller images
+            transforms.CenterCrop(512),  # Then crop to 512x512
         ])
 
         self.compression_formats = ['JPEG', 'WebP', 'AVIF', 'JXL']
@@ -64,7 +70,7 @@ class CompressionArtifactPredictor:
         """Predict compression quality levels for all formats."""
         img_tensor = self.preprocess(image).unsqueeze(0).to(self.device)
 
-        # SIMPLE FULL PRECISION INFERENCE - NO AUTOCAST
+        # FIXED: Full precision, no autocast
         with torch.no_grad():
             predictions = self.model(img_tensor).squeeze(0).cpu().numpy()
 
@@ -99,7 +105,6 @@ class CompressionArtifactPredictor:
 
         return results
 
-
 # ==================== GRADIO UI ====================
 def create_ui():
     predictor = CompressionArtifactPredictor()
@@ -112,25 +117,27 @@ def create_ui():
             image = Image.fromarray(image)
         
         image = image.convert('RGB')
+        print(f"Processing image of size: {image.size}")  # Debug log
+        
         results = predictor.predict(image)
         
-        # Generate HTML table for results
+        # FIXED: Dark mode compatible using CSS variables
         html_results = """
         <table style='width:100%; border-collapse: collapse; font-family: inherit;'>
-            <tr style='background: #f5f5f5;'>
-                <th style='padding:12px; text-align:left; border-bottom: 2px solid #ddd;'>Format</th>
-                <th style='padding:12px; text-align:center; border-bottom: 2px solid #ddd;'>Quality</th>
-                <th style='padding:12px; text-align:center; border-bottom: 2px solid #ddd;'>Assessment</th>
-                <th style='padding:12px; text-align:center; border-bottom: 2px solid #ddd;'>Accuracy</th>
+            <tr style='background: var(--block-label-background-fill, #f5f5f5);'>
+                <th style='padding:12px; text-align:left; border-bottom: 2px solid var(--border-color-primary, #ddd);'>Format</th>
+                <th style='padding:12px; text-align:center; border-bottom: 2px solid var(--border-color-primary, #ddd);'>Quality</th>
+                <th style='padding:12px; text-align:center; border-bottom: 2px solid var(--border-color-primary, #ddd);'>Assessment</th>
+                <th style='padding:12px; text-align:center; border-bottom: 2px solid var(--border-color-primary, #ddd);'>Accuracy</th>
             </tr>
         """
         
         for fmt, data in results.items():
             html_results += f"""
-                <tr style='border-bottom: 1px solid #eee;'>
+                <tr style='border-bottom: 1px solid var(--border-color-primary, #eee);'>
                     <td style='padding:12px; font-weight:500;'>{data['indicator']} {fmt}</td>
                     <td style='padding:12px; text-align:center;'><strong>{data['quality_score']}/100</strong></td>
-                    <td style='padding:12px; text-align:center;'>{data['category']}<br><small style='color:#666;'>{data['desc']}</small></td>
+                    <td style='padding:12px; text-align:center;'>{data['category']}<br><small style='color: var(--body-text-color-subdued, #666);'>{data['desc']}</small></td>
                     <td style='padding:12px; text-align:center;'>{data['accuracy']}%</td>
                 </tr>
             """
@@ -181,7 +188,6 @@ def create_ui():
                 analyze_button = gr.Button("🔍 Analyze Image Quality", variant="primary", size="lg")
             
             with gr.Column():
-                # USE HTML COMPONENT INSTEAD OF LABEL
                 results_output = gr.HTML(
                     label="Format-Specific Quality Scores"
                 )