add: cropping

src/streamlit_app.py

@@ -2,6 +2,26 @@ import streamlit as st
 import fitz
 from PIL import Image
 import io
+from streamlit_cropper import st_cropper
+
+
+def get_padded_image(image, cropped_image):
+    base_w, base_h = image.size
+    crop_w, crop_h = cropped_image.size
+    # Match modes to avoid paste issues
+    crop_img = (
+        cropped_image.convert(image.mode)
+        if cropped_image.mode != image.mode
+        else cropped_image
+    )
+    # Create white background canvas matching the original image size
+    padded_image = Image.new(image.mode, (base_w, base_h), color="white")
+    # Center the cropped image on the canvas
+    paste_x = max(0, (base_w - crop_w) // 2)
+    paste_y = max(0, (base_h - crop_h) // 2)
+    padded_image.paste(crop_img, (paste_x, paste_y))
+    return padded_image
+
 
 st.header("Line Art Data Annotation App")
 uploaded_pdf = st.sidebar.file_uploader("Upload a PDF", type=["pdf"])
@@ -9,10 +29,12 @@ uploaded_pdf = st.sidebar.file_uploader("Upload a PDF", type=["pdf"])
 if uploaded_pdf:
     data = uploaded_pdf.read()
     doc = fitz.open(stream=data, filetype="pdf")
-    
+
     # Initialize page index in session state
     if "page_idx" not in st.session_state:
         st.session_state.page_idx = 0
+    if "cropped_images" not in st.session_state:
+        st.session_state.cropped_images = []
 
     total_pages = doc.page_count
     page_idx = st.session_state.page_idx % total_pages
@@ -23,7 +45,10 @@ if uploaded_pdf:
             st.session_state.page_idx = (page_idx - 1) % total_pages
             st.rerun()
     with col_caption:
-        st.markdown(f"<center>Page {page_idx}/{total_pages - 1}</center>", unsafe_allow_html=True)
+        st.markdown(
+            f"<center>Page {page_idx}/{total_pages - 1}</center>",
+            unsafe_allow_html=True,
+        )
     with col_next:
         if st.button("\>"):
             st.session_state.page_idx = (page_idx + 1) % total_pages
@@ -32,5 +57,11 @@ if uploaded_pdf:
     # Render image after controls
     page = doc.load_page(page_idx)
     pix = page.get_pixmap(dpi=200)
-    img = Image.open(io.BytesIO(pix.tobytes("png")))
-    st.image(img)
+    image = Image.open(io.BytesIO(pix.tobytes("png")))
+    cropped_image = st_cropper(image, realtime_update=True)
+    st.image(cropped_image)
+    if st.button("Save"):
+        padded_image = get_padded_image(image, cropped_image)
+        st.session_state.cropped_images.append(padded_image)
+        print(f"{len(st.session_state.cropped_images)=}")
+        print(f"{st.session_state.cropped_images[-1].size=}")