Update app.py

app.py

@@ -1,16 +1,14 @@
-import numpy as np
-import streamlit as st
+import gradio as gr
 import requests
 from PIL import Image
 from io import BytesIO
-from copy import deepcopy
-from streamlit_drawable_canvas import st_canvas
 
-# Hàm để tải ảnh và mask lên API Magic Studio
+# Hàm xử lý gửi ảnh và mask tới API Magic Studio
 def process_inpaint_via_api(image, mask):
     url = "https://api.magicstudio.com/magiceraser/erase"
+    
     headers = {
-        'accessToken': 'eyJ0eXAiOiJKV1QiLCJhbGciOiJIUzI1NiJ9.eyJjbGllbnRfaWQiOiIyRGZQSm1sY1BLalBuaGh4ajBxSVpmQktfbS12RFhTc1NMRFNOX1gtZlVVIiwiZXhwIjoxNzI3MTY0MTcyLCJhcHBfbmFtZSI6IjE3NTMwODQiLCJtZXRhZGF0YSI6bnVsbCwiZ3JhbnRfdHlwZSI6ImNyZWQifQ.24oHGH9ial_gqeYrpUbYP68MwYly2ArzXbhbE1hGzQo'
+        'accessToken': 'eyJ0eXAiOiJKV1QiLCJhbGciOiJIUzI1NiJ9.eyJjbGllbnRfaWQiOiIyRGZQSm1sY1BLalBuaGh4ajBxSVpmQktfbS12RFhTc1NMRFNOX1gtZlVVIiwiZXhwIjoxNzI3MTY0MTcyLCJhcHBfbmFtZSI6IjE3NTMwODQiLCJtZXRhbnRfdHlwZSI6ImNyZWQifQ.24oHGH9ial_gqeYrpUbYP68MwYly2ArzXbhbE1hGzQo'
     }
     
     # Chuyển đổi ảnh và mask thành định dạng PNG để gửi tới API
@@ -36,77 +34,38 @@ def process_inpaint_via_api(image, mask):
         return Image.open(BytesIO(response.content))
     else:
         # Nếu có lỗi, hiển thị lỗi
-        st.error(f"API Error: {response.status_code} - {response.text}")
-        return None
-
-# Thiết lập giao diện Streamlit
-st.title("AI Photo Object Removal")
-
-# Tải ảnh từ người dùng
-uploaded_file = st.file_uploader("Choose image", accept_multiple_files=False, type=["png", "jpg", "jpeg"])
-
-if uploaded_file is not None:
-    # Đọc ảnh từ người dùng
-    img_input = Image.open(uploaded_file).convert("RGBA")
+        return f"API Error: {response.status_code} - {response.text}"
 
-    # 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))
+# Hàm xử lý chính trong Gradio
+def process(image, mask):
+    img = Image.fromarray(image)
+    mask_img = Image.fromarray(mask)
     
-    # Lựa chọn kích thước cọ vẽ
-    stroke_width = st.slider("Brush size", 1, 100, 50)
-
-    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
+    # Gọi API với ảnh và mask
+    result = process_inpaint_via_api(img, mask_img)
     
-    # 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)))
+    # Nếu kết quả trả về là lỗi, hiển thị thông báo
+    if isinstance(result, str):
+        return None, result
     
-    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"
-    )
+    # Trả về ảnh đã xử lý
+    return result, "Processing complete."
+
+# Tạo giao diện Gradio
+with gr.Blocks() as demo:
+    gr.Markdown("# AI Object Removal using Magic Studio API")
     
-    # 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)
+    # Khu vực tải lên ảnh và vẽ mask
+    with gr.Row():
+        image_input = gr.Image(source="upload", tool="sketch", type="numpy", label="Upload an image and draw the mask")
 
-        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
+    # Khu vực hiển thị kết quả
+    result_output = gr.Image(label="Output Image")
+    output_message = gr.Textbox(label="Message", lines=2)
+    
+    # Nút xử lý và kết nối các thành phần với hàm process
+    process_button = gr.Button("Submit")
+    process_button.click(process, inputs=[image_input, image_input], outputs=[result_output, output_message])
 
-            # 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)
+# Khởi chạy ứng dụng Gradio
+demo.launch()