Update app.py

app.py

@@ -1,66 +1,112 @@
-import gradio as gr
+import numpy as np
+import streamlit as st
 import requests
 from PIL import Image
-import numpy as np
+from io import BytesIO
+from copy import deepcopy
+from streamlit_drawable_canvas import st_canvas
 
-# Hàm gọi API Magic Studio để thực hiện xóa đối tượng
-def erase_object(image, mask, access_token):
+# Hàm để tải ảnh và mask lên API Magic Studio
+def process_inpaint_via_api(image, mask):
     url = "https://api.magicstudio.com/magiceraser/erase"
+    headers = {
+        'accessToken': 'eyJ0eXAiOiJKV1QiLCJhbGciOiJIUzI1NiJ9.eyJjbGllbnRfaWQiOiIyRGZQSm1sY1BLalBuaGh4ajBxSVpmQktfbS12RFhTc1NMRFNOX1gtZlVVIiwiZXhwIjoxNzI3MTY0MTcyLCJhcHBfbmFtZSI6IjE3NTMwODQiLCJtZXRhZGF0YSI6bnVsbCwiZ3JhbnRfdHlwZSI6ImNyZWQifQ.24oHGH9ial_gqeYrpUbYP68MwYly2ArzXbhbE1hGzQo'
+    }
     
-    # Chuyển đổi hình ảnh và mask thành định dạng thích hợp để gửi qua API
-    image_file = ('image_file', ('image.jpg', image, 'image/jpeg'))
+    # Chuyển đổi ảnh và mask thành định dạng PNG để gửi tới API
+    image_io = BytesIO()
+    mask_io = BytesIO()
     
-    # Mask là một hình ảnh do người dùng vẽ trực tiếp từ Gradio Sketchpad
-    mask_file = ('mask_file', ('mask.png', mask, 'image/png'))
+    image.save(image_io, format="PNG")
+    mask.save(mask_io, format="PNG")
     
-    # Payload và headers
-    payload = {'filename': 'image.jpg'}
-    headers = {'accessToken': access_token}
+    image_io.seek(0)
+    mask_io.seek(0)
     
-    # Gửi yêu cầu POST đến API với hình ảnh và mask
-    response = requests.post(url, headers=headers, data=payload, files=[image_file, mask_file])
+    files = [
+        ('image_file', ('image.png', image_io, 'image/png')),
+        ('mask_file', ('mask.png', mask_io, 'image/png'))
+    ]
+    
+    # Gửi yêu cầu POST tới API
+    response = requests.post(url, headers=headers, files=files)
     
     if response.status_code == 200:
-        # Trả về hình ảnh đã xóa đối tượng nếu thành công
-        return response.json()
+        # Nếu API trả về thành công, xử lý và trả về ảnh kết quả
+        return Image.open(BytesIO(response.content))
     else:
-        # Trả về thông báo lỗi nếu thất bại
-        return f"Error: {response.status_code} - {response.text}"
-
-# Hàm xử lý cho Gradio
-def process(image, mask):
-    # Chuyển đổi hình ảnh từ Gradio thành định dạng 'bytes'
-    image_bytes = image.convert("RGB").tobytes()
-    mask_bytes = mask.convert("RGB").tobytes()
-    
-    # Gọi API với hình ảnh và mask (token được cố định trong code)
-    access_token = 'eyJ0eXAiOiJKV1QiLCJhbGciOiJIUzI1NiJ9.eyJjbGllbnRfaWQiOiIyRGZQSm1sY1BLalBuaGh4ajBxSVpmQktfbS12RFhTc1NMRFNOX1gtZlVVIiwiZXhwIjoxNzI3MTY0MTcyLCJhcHBfbmFtZSI6IjE3NTMwODQiLCJtZXRhZGF0YSI6bnVsbCwiZ3JhbnRfdHlwZSI6ImNyZWQifQ.24oHGH9ial_gqeYrpUbYP68MwYly2ArzXbhbE1hGzQo'
-    
-    result = erase_object(image_bytes, mask_bytes, access_token)
-    
-    return result
+        # Nếu có lỗi, hiển thị lỗi
+        st.error(f"API Error: {response.status_code} - {response.text}")
+        return None
 
-# Tạo giao diện Gradio
-with gr.Blocks() as demo:
-    gr.Markdown("# Magic Eraser - Xóa đối tượng khỏi ảnh")
+# Thiết lập giao diện Streamlit
+st.title("AI Photo Object Removal")
 
-    # Input cho ảnh và Sketchpad để vẽ mask
-    with gr.Row():
-        image_input = gr.Image(type="pil", label="Tải ảnh lên")
-        mask_input = gr.Sketchpad(shape=(512, 512), label="Vẽ mask để xóa đối tượng")
+# Tải ảnh từ người dùng
+uploaded_file = st.file_uploader("Choose image", accept_multiple_files=False, type=["png", "jpg", "jpeg"])
 
-    # Nút thực hiện xử lý xóa đối tượng
-    submit_button = gr.Button("Xóa đối tượng")
+if uploaded_file is not None:
+    # Đọc ảnh từ người dùng
+    img_input = Image.open(uploaded_file).convert("RGBA")
 
-    # Kết quả trả về là ảnh đã xóa đối tượng
-    result_output = gr.Image(type="pil", label="Kết quả")
+    # Resize ảnh để phù hợp với giao diện
+    max_size = 2000
+    img_width, img_height = img_input.size
+    if img_width > max_size or img_height > max_size:
+        if img_width > img_height:
+            new_width = max_size
+            new_height = int((max_size / img_width) * img_height)
+        else:
+            new_height = max_size
+            new_width = int((max_size / img_height) * img_width)
+        img_input = img_input.resize((new_width, new_height))
+    
+    # Lựa chọn kích thước cọ vẽ
+    stroke_width = st.slider("Brush size", 1, 100, 50)
 
-    # Kết nối các thành phần với nhau
-    submit_button.click(
-        process,
-        inputs=[image_input, mask_input],
-        outputs=result_output
+    st.write("**Now draw (brush) the part of image that you want to remove.**")
+    
+    # Canvas để người dùng vẽ mask trên ��nh
+    canvas_bg = deepcopy(img_input)
+    aspect_ratio = canvas_bg.width / canvas_bg.height
+    streamlit_width = 720
+    
+    # Resize ảnh nền để hiển thị đúng kích thước trên Streamlit
+    if canvas_bg.width > streamlit_width:
+        canvas_bg = canvas_bg.resize((streamlit_width, int(streamlit_width / aspect_ratio)))
+    
+    canvas_result = st_canvas(
+        stroke_color="rgba(255, 0, 255, 1)",  # Màu của cọ vẽ mask
+        stroke_width=stroke_width,  # Độ dày cọ vẽ
+        background_image=canvas_bg,  # Ảnh nền
+        width=canvas_bg.width,
+        height=canvas_bg.height,
+        drawing_mode="freedraw",  # Chế độ vẽ tự do
+        key="canvas"
     )
+    
+    # Xử lý khi người dùng đã vẽ mask và nhấn nút Submit
+    if canvas_result.image_data is not None:
+        # Chuyển đổi dữ liệu từ canvas thành ảnh
+        im = np.array(Image.fromarray(canvas_result.image_data.astype(np.uint8)).resize(img_input.size))
+        background = np.where(
+            (im[:, :, 0] == 0) & 
+            (im[:, :, 1] == 0) & 
+            (im[:, :, 2] == 0)
+        )
+        drawing = np.where(
+            (im[:, :, 0] == 255) & 
+            (im[:, :, 1] == 0) & 
+            (im[:, :, 2] == 255)
+        )
+        im[background] = [0, 0, 0, 255]  # Đặt nền thành đen
+        im[drawing] = [0, 0, 0, 0]  # Đặt mask thành trong suốt (alpha)
+
+        if st.button('Submit'):
+            with st.spinner("AI is doing the magic!"):
+                img_output = process_inpaint_via_api(img_input, Image.fromarray(im))  # Gọi API xử lý inpainting
 
-# Khởi chạy ứng dụng Gradio
-demo.launch()
+            # Hiển thị ảnh kết quả sau khi xử lý
+            if img_output is not None:
+                st.write("AI has finished the job!")
+                st.image(img_output)