Add main space test file

app.py

@@ -28,9 +28,12 @@ except Exception as e:
 
 def load_and_preprocess_image(image_path):
     image = load_rgb(image_path)
-    padded_image, pads = pad(image, factor=32, border=cv2.BORDER_CONSTANT)
-    x = torch.unsqueeze(tensor_from_rgb_image(image), 0)  # Use original image, not padded
-    return x, image, pads
+    # Pad the image to be divisible by 32
+    padded_image, pads = pad(image, factor=32, border=cv2.BORDER_CONSTANT)  
+    transform = albu.Compose([albu.Normalize(p=1)], p=1)
+    x = transform(image=padded_image)["image"]  # Use padded image here
+    x = torch.unsqueeze(tensor_from_rgb_image(x), 0)
+    return x, image, pads  # Return original image and padding info for later unpadding
 
 def segment_cloth(image_tensor, model, pads):
     with torch.no_grad():
@@ -56,19 +59,26 @@ def image_segmentation_and_inpainting(image, prompt="Chinese Red and Golder Armo
         temp_image_path = "temp_image.jpg"
         pil_image.save(temp_image_path)
 
-        x, image, pads = load_and_preprocess_image(temp_image_path)
+        x, original_image, pads = load_and_preprocess_image(temp_image_path) 
         mask = segment_cloth(x, model, pads)
 
+        # Ensure mask has same dimensions as original image before saving
+        mask = unpad(mask, pads)
+
         mask_path = "temp_mask.jpg"
         plt.imsave(mask_path, mask, cmap='gray')
 
         output_image = perform_inpainting(temp_image_path, mask_path, prompt)
-        output_image = resize_and_upscale(output_image, 1280, 720)  # You can adjust the size
+
+        # Resize the output image to match the original image's dimensions
+        output_image = resize_and_upscale(output_image, original_image.shape[1], original_image.shape[0])
 
         return output_image
     except Exception as e:
         raise gr.Error(f"Error processing image: {e}")
 
+
+
 with gr.Blocks() as demo:
     gr.Markdown("# Cloth Image Segmentation and Inpainting")
     with gr.Row():