Update app.py

app.py

@@ -24,11 +24,13 @@ print(f"Model loaded on {device}")
 OUTPUT_DIR = "output_images"
 os.makedirs(OUTPUT_DIR, exist_ok=True)
 
+# Define fixed size for all images
+FIXED_SIZE = (512, 512)
+
 # Resize the input image for model compatibility
-def resize_image(image):
+def resize_image(image, size=FIXED_SIZE):
     image = image.convert('RGB')
-    model_input_size = (1024, 1024)
-    image = image.resize(model_input_size, Image.BILINEAR)
+    image = image.resize(size, Image.LANCZOS)
     return image
 
 # Background removal process
@@ -41,8 +43,10 @@ def process(image, progress=gr.Progress()):
     # Prepare the input
     progress(0.1, desc="Preparing image...")
     orig_image = Image.fromarray(image)
-    w, h = orig_im_size = orig_image.size
-    image = resize_image(orig_image)
+    # Resize input image to fixed size
+    orig_image = resize_image(orig_image, FIXED_SIZE)
+    w, h = FIXED_SIZE
+    image = orig_image
     im_np = np.array(image)
     im_tensor = torch.tensor(im_np, dtype=torch.float32).permute(2, 0, 1)
     im_tensor = torch.unsqueeze(im_tensor, 0)
@@ -86,11 +90,8 @@ def process(image, progress=gr.Progress()):
     
     progress(1.0, desc="Done!")
     
-    # Return the processed image and the file for download
     return output_array, gr.update(visible=True, value=filepath, interactive=True)
 
-# ... previous code remains the same until the description ...
-
 # Gradio interface setup
 title = "Background Removal Tool"
 description = """
@@ -105,7 +106,7 @@ description = """
         box-shadow: 0 0 20px rgba(0, 255, 255, 0.2);
     }
     .title-text {
-        color: #ffffff;
+        color: #ff00de;
         font-family: 'Orbitron', sans-serif;
         font-size: 2.5em;
         margin: 20px 0;
@@ -113,7 +114,7 @@ description = """
         animation: title-pulse 2s infinite alternate;
     }
     .subtitle-text {
-        color: #ffffff;
+        color: #00ffff;
         font-family: 'Roboto Mono', monospace;
         font-size: 1.2em;
         margin-top: 10px;
@@ -167,6 +168,9 @@ with gr.Blocks(css="""
 
     /* Input/Output areas */
     .input-image, .output-image {
+        width: 512px !important;
+        height: 512px !important;
+        object-fit: contain !important;
         background: rgba(18, 18, 56, 0.7) !important;
         border: 2px solid var(--neon-cyan) !important;
         border-radius: 12px !important;
@@ -174,6 +178,12 @@ with gr.Blocks(css="""
         overflow: hidden !important;
     }
 
+    .input-image img, .output-image img {
+        width: 100% !important;
+        height: 100% !important;
+        object-fit: contain !important;
+    }
+
     .input-image:hover, .output-image:hover {
         box-shadow: 0 0 20px rgba(0, 255, 255, 0.4) !important;
         transform: translateY(-2px) !important;
@@ -232,7 +242,9 @@ with gr.Blocks(css="""
                 type="numpy",
                 label="Upload Your Image",
                 elem_id="input-image",
-                elem_classes="input-image"
+                elem_classes="input-image",
+                height=512,
+                width=512
             )
         
         with gr.Column(scale=1):
@@ -240,7 +252,9 @@ with gr.Blocks(css="""
                 type="numpy",
                 label="Result",
                 elem_id="output-image",
-                elem_classes="output-image"
+                elem_classes="output-image",
+                height=512,
+                width=512
             )
             
             with gr.Row(elem_classes="download-container"):