push app

app.py

@@ -21,6 +21,18 @@ st.set_page_config(page_title="Inpainting Demo", layout="wide")
 FASTSAM_CHECKPOINT = "FastSAM-x.pt"  # file name of the FastSAM model weights
 SD_MODEL_ID = "runwayml/stable-diffusion-inpainting"  # HF Hub model for SD Inpainting v1.5
 
+# Helper function: center crop and resize to 512x512
+def crop_resize_image(image, size=512):
+    width, height = image.size
+    if width != height:
+        min_dim = min(width, height)
+        left = (width - min_dim) // 2
+        top = (height - min_dim) // 2
+        right = left + min_dim
+        bottom = top + min_dim
+        image = image.crop((left, top, right, bottom))
+    return image.resize((size, size))
+
 # Ensure FastSAM model weights are available (download if not present)
 if not os.path.exists(FASTSAM_CHECKPOINT):
     # Download FastSAM weights (if not already in the repo)
@@ -68,22 +80,27 @@ canvas = st.empty()
 picture = st.camera_input("Take a picture")
 
 if picture is not None:
-    # When an image is captured, display it and allow point selection
-    img = Image.open(picture)  # read image as PIL
+    # Open, crop & resize the captured image to 512x512
+    img = Image.open(picture)
+    img = crop_resize_image(img, size=512)
     canvas.image(img, caption="Captured Image", use_container_width=True)
     
-    # Let user click a point on the image. This returns a dict with 'x' and 'y'.
-    coords = streamlit_image_coordinates(img, key="click_img")
-    if coords:
-        # If a point was clicked, mark it on the image for user feedback
-        cx, cy = int(coords['x']), int(coords['y'])
-        # Draw a small red circle on the image copy to show selected point
-        img_with_dot = img.copy()
-        draw = ImageDraw.Draw(img_with_dot)
-        draw.ellipse((cx-5, cy-5, cx+5, cy+5), fill='red')
-        canvas.image(img_with_dot, caption=f"Selected Point: ({cx}, {cy})", use_container_width=True)
-    else:
-        cx = cy = None
+    # Use the canvas container for all related UI elements
+    with canvas.container():
+        st.image(img, caption="Captured (512x512) Image", use_container_width=True)
+        
+        # Place the interactive component inside the same container
+        coords = streamlit_image_coordinates(img, key="click_img")
+        
+        if coords:
+            cx, cy = int(coords['x']), int(coords['y'])
+            # Draw a red circle on the image to indicate the selected point
+            img_with_dot = img.copy()
+            draw = ImageDraw.Draw(img_with_dot)
+            draw.ellipse((cx-5, cy-5, cx+5, cy+5), fill='red')
+            st.image(img_with_dot, caption=f"Selected Point: ({cx}, {cy})", use_container_width=True)
+        else:
+            cx = cy = None
 
     # Prompt input for inpainting
     prompt = st.text_input("Prompt for inpainting (describe what should replace the selected area):")
@@ -92,21 +109,14 @@ if picture is not None:
     if coords and prompt:
         cx, cy = int(coords['x']), int(coords['y'])
         st.write("Generating mask with FastSAM...")
-        # Run FastSAM segmentation with the selected point as prompt
-        # Using the Ultralytics API: points=[[x,y]] and labels=[1] for a positive point prompt
         results = fastsam_model(img, points=[[cx, cy]], labels=[1])
-        # The results object holds masks; extract the first mask (closest object to the point)
-        mask_data = results[0].masks.data[0]  # mask tensor (H x W)
-        mask_array = mask_data.cpu().numpy()  # convert to numpy array
-        # Create a PIL Image for the mask: convert 1.0 to 255 (white), 0.0 to 0 (black)
+        mask_data = results[0].masks.data[0]
+        mask_array = mask_data.cpu().numpy()
         mask_image = Image.fromarray((mask_array * 255).astype(np.uint8))
         
-        # For debugging, we can display the mask – uncomment if needed
-        # st.image(mask_image, caption="Segmentation Mask", use_column_width=True)
-        
         st.write("Running Stable Diffusion Inpainting...")
-        # Run the Stable Diffusion inpainting pipeline
         result = sd_pipe(prompt=prompt, image=img, mask_image=mask_image).images[0]
         
-        # Display the final inpainted image
-        canvas.image(result, caption="Inpainted Image", use_container_width=True)
+        # Finally, update the same canvas with the inpainted image
+        with canvas.container():
+            st.image(result, caption="Inpainted Image", use_container_width=True)