Simple refactor for readibility

app.py

@@ -10,6 +10,12 @@ import glob
 import base64
 from io import BytesIO
 
+# Constants
+MODEL_CHECKPOINT = 'net_g_150000.pth'
+REQUIRED_IMAGE_SIZE = (130, 130)
+WINDOW_SIZE = 16
+UPSCALE_FACTOR = 4
+
 
 def to_2tuple(x):
     """Convert input to tuple of length 2."""
@@ -685,13 +691,10 @@ model = HAT(
 )
 
 # Load the fine-tuned weights
-checkpoint = torch.load('net_g_150000.pth', map_location=device)
-if 'params_ema' in checkpoint:
-    model.load_state_dict(checkpoint['params_ema'])
-elif 'params' in checkpoint:
-    model.load_state_dict(checkpoint['params'])
-else:
-    model.load_state_dict(checkpoint)
+checkpoint = torch.load(MODEL_CHECKPOINT, map_location=device)
+# Try different checkpoint formats
+state_dict = checkpoint.get('params_ema') or checkpoint.get('params') or checkpoint
+model.load_state_dict(state_dict)
 
 model.to(device)
 model.eval()
@@ -706,8 +709,8 @@ def upscale_image(image):
     h, w = img_tensor.shape[2], img_tensor.shape[3]
 
     # Pad if necessary
-    pad_h = (16 - h % 16) % 16
-    pad_w = (16 - w % 16) % 16
+    pad_h = (WINDOW_SIZE - h % WINDOW_SIZE) % WINDOW_SIZE
+    pad_w = (WINDOW_SIZE - w % WINDOW_SIZE) % WINDOW_SIZE
 
     if pad_h > 0 or pad_w > 0:
         img_tensor = torch.nn.functional.pad(img_tensor, (0, pad_w, 0, pad_h), mode='reflect')
@@ -717,7 +720,7 @@ def upscale_image(image):
 
     # Remove padding if it was added
     if pad_h > 0 or pad_w > 0:
-        output = output[:, :, :h*4, :w*4]
+        output = output[:, :, :h*UPSCALE_FACTOR, :w*UPSCALE_FACTOR]
 
     # Convert back to PIL image
     output_np = output.squeeze(0).permute(1, 2, 0).cpu().numpy()
@@ -736,13 +739,14 @@ def get_sample_images():
 
 # Gradio interface using Blocks for better layout control
 def validate_image_size(image):
-    """Validate that the image is exactly 130x130 pixels"""
+    """Validate that the image is exactly the required size"""
     if image is None:
         return False, "No image provided"
 
     width, height = image.size
-    if width != 130 or height != 130:
-        return False, f"Image must be exactly 130x130 pixels. Your image is {width}x{height} pixels."
+    req_width, req_height = REQUIRED_IMAGE_SIZE
+    if width != req_width or height != req_height:
+        return False, f"Image must be exactly {req_width}x{req_height} pixels. Your image is {width}x{height} pixels."
 
     return True, "Valid image size"
 
@@ -811,11 +815,16 @@ def generate_css():
 }
 """
 
-    # Add background images for each sample
+    # Add background images for each sample (only if samples exist)
     sample_images = get_sample_images()
-    for i, img_path in enumerate(sample_images):
-        base64_img = image_to_base64(img_path)
-        base_css += f"#sample_btn_{i} {{ background-image: url('{base64_img}'); }}\n"
+    if sample_images:
+        for i, img_path in enumerate(sample_images):
+            try:
+                base64_img = image_to_base64(img_path)
+                base_css += f"#sample_btn_{i} {{ background-image: url('{base64_img}'); }}\n"
+            except Exception:
+                # Skip invalid images
+                continue
 
     return base_css
 
@@ -823,8 +832,8 @@ css = generate_css()
 
 with gr.Blocks(css=css, title="HATSAT - Super-Resolution for Satellite Images") as iface:
     gr.Markdown("# HATSAT - Super-Resolution for Satellite Images")
-    gr.Markdown("Upload a satellite image or select a sample to enhance its resolution by 4x.")
-    gr.Markdown("⚠️ **Important**: Images must be exactly **130x130 pixels** for the model to work properly.")
+    gr.Markdown(f"Upload a satellite image or select a sample to enhance its resolution by {UPSCALE_FACTOR}x.")
+    gr.Markdown(f"⚠️ **Important**: Images must be exactly **{REQUIRED_IMAGE_SIZE[0]}x{REQUIRED_IMAGE_SIZE[1]} pixels** for the model to work properly.")
 
     # Acknowledgments section
     with gr.Accordion("Acknowledgments", open=False):
@@ -858,7 +867,7 @@ with gr.Blocks(css=css, title="HATSAT - Super-Resolution for Satellite Images")
     with gr.Row():
         input_image = gr.Image(
             type="pil",
-            label="Input Image (must be 130x130 pixels)",
+            label=f"Input Image (must be {REQUIRED_IMAGE_SIZE[0]}x{REQUIRED_IMAGE_SIZE[1]} pixels)",
             elem_classes="image-container",
             sources=["upload"],
             height=500,
@@ -867,7 +876,7 @@ with gr.Blocks(css=css, title="HATSAT - Super-Resolution for Satellite Images")
 
         output_image = gr.Image(
             type="pil",
-            label="Enhanced Output (4x)",
+            label=f"Enhanced Output ({UPSCALE_FACTOR}x)",
             elem_classes="image-container",
             interactive=False,
             height=500,