Update app.py

app.py

@@ -24,17 +24,26 @@ print(f"Model loaded on {device}")
 OUTPUT_DIR = "output_images"
 os.makedirs(OUTPUT_DIR, exist_ok=True)
 
-# Define different sizes for different devices
-DESKTOP_SIZE = (512, 512)
-MOBILE_SIZE = (300, 300)
-
-# Resize the input image for model compatibility
-def resize_image(image, size=DESKTOP_SIZE):
-    image = image.convert('RGB')
-    image = image.resize(size, Image.LANCZOS)
+def resize_image(image, max_size=1024):
+    """Resize image while maintaining aspect ratio and quality"""
+    # Get original size
+    width, height = image.size
+    
+    # Calculate aspect ratio
+    aspect_ratio = width / height
+    
+    # Only resize if the image is larger than max_size in either dimension
+    if width > max_size or height > max_size:
+        if width > height:
+            new_width = max_size
+            new_height = int(max_size / aspect_ratio)
+        else:
+            new_height = max_size
+            new_width = int(max_size * aspect_ratio)
+        image = image.resize((new_width, new_height), Image.LANCZOS)
+    
     return image
 
-# Background removal process
 def process(image, progress=gr.Progress()):
     if image is None:
         return None, gr.update(visible=False)
@@ -44,11 +53,13 @@ def process(image, progress=gr.Progress()):
     # Prepare the input
     progress(0.1, desc="Preparing image...")
     orig_image = Image.fromarray(image)
-    # Resize input image to fixed size
-    orig_image = resize_image(orig_image, DESKTOP_SIZE)
-    w, h = DESKTOP_SIZE
-    image = orig_image
-    im_np = np.array(image)
+    original_size = orig_image.size
+    
+    # Resize only if needed for processing
+    process_image = resize_image(orig_image)
+    w, h = process_image.size
+    
+    im_np = np.array(process_image)
     im_tensor = torch.tensor(im_np, dtype=torch.float32).permute(2, 0, 1)
     im_tensor = torch.unsqueeze(im_tensor, 0)
     im_tensor = torch.divide(im_tensor, 255.0)
@@ -73,6 +84,10 @@ def process(image, progress=gr.Progress()):
     result_array = (result * 255).cpu().data.numpy().astype(np.uint8)
     pil_mask = Image.fromarray(np.squeeze(result_array))
     
+    # Resize mask back to original size if needed
+    if pil_mask.size != original_size:
+        pil_mask = pil_mask.resize(original_size, Image.LANCZOS)
+    
     # Add the mask as alpha channel to the original image
     new_im = orig_image.copy()
     new_im.putalpha(pil_mask)
@@ -83,8 +98,8 @@ def process(image, progress=gr.Progress()):
     filename = f"background_removed_{unique_id}.png"
     filepath = os.path.join(OUTPUT_DIR, filename)
     
-    # Save the processed image
-    new_im.save(filepath, format='PNG')
+    # Save the processed image in original resolution
+    new_im.save(filepath, format='PNG', quality=100)
     
     # Convert to numpy array for display
     output_array = np.array(new_im.convert("RGBA"))
@@ -178,9 +193,9 @@ with gr.Blocks(css="""
     /* Input/Output areas with responsive sizing */
     .input-image, .output-image {
         width: 100% !important;
-        max-width: 512px !important;
+        max-width: 800px !important;
         height: auto !important;
-        aspect-ratio: 1/1 !important;
+        min-height: 300px !important;
         object-fit: contain !important;
         background: rgba(18, 18, 56, 0.7) !important;
         border: 2px solid var(--neon-cyan) !important;
@@ -191,9 +206,10 @@ with gr.Blocks(css="""
     }
 
     .input-image img, .output-image img {
-        width: 100% !important;
-        height: 100% !important;
+        max-width: 100% !important;
+        max-height: 800px !important;
         object-fit: contain !important;
+        margin: auto !important;
     }
 
     /* Responsive columns */
@@ -236,7 +252,11 @@ with gr.Blocks(css="""
     /* Responsive layout */
     @media (max-width: 768px) {
         .input-image, .output-image {
-            max-width: 300px !important;
+            min-height: 200px !important;
+        }
+        
+        .input-image img, .output-image img {
+            max-height: 500px !important;
         }
         
         label {