Spaces:
Sleeping
Sleeping
| !pip install git+https://github.com/facebookresearch/segment-anything.git | |
| !pip install opencv-python pillow matplotlib | |
| !pip uninstall -y diffusers transformers | |
| !pip install diffusers==0.30.3 transformers==4.44.2 accelerate | |
| !pip install gradio | |
| import gradio as gr | |
| import numpy as np | |
| import torch | |
| from PIL import Image, ImageDraw | |
| import cv2 | |
| from segment_anything import sam_model_registry, SamPredictor | |
| from diffusers import StableDiffusionXLInpaintPipeline | |
| import os | |
| # ============================================================ | |
| # SAM SMART FOREGROUND SELECTOR - OFFLINE MODEL VERSION | |
| # ============================================================ | |
| # Global state | |
| global_state = { | |
| "sam_predictor": None, | |
| "original_image": None, | |
| "current_mask": None, | |
| "rgba_image": None, | |
| "image_set": False, | |
| "box_points": [], | |
| "positive_points": [], | |
| "negative_points": [], | |
| "auto_masks": [] | |
| } | |
| # ------------------------------------------------------------ | |
| # 🧩 Load SAM Model (Offline from /models/) | |
| # ------------------------------------------------------------ | |
| def load_sam_model(sam_checkpoint="models/sam_vit_b_01ec64.pth"): | |
| """Load SAM model from local folder models/""" | |
| if global_state["sam_predictor"] is None: | |
| if not os.path.exists(sam_checkpoint): | |
| raise FileNotFoundError( | |
| f"❌ File model tidak ditemukan: {sam_checkpoint}\n" | |
| f"Pastikan kamu sudah upload 'sam_vit_b_01ec64.pth' ke folder /models/" | |
| ) | |
| print(f"✅ Loading SAM model from {sam_checkpoint}...") | |
| sam = sam_model_registry["vit_b"](checkpoint=sam_checkpoint) | |
| device = "cuda" if torch.cuda.is_available() else "cpu" | |
| sam.to(device=device) | |
| predictor = SamPredictor(sam) | |
| global_state["sam_predictor"] = predictor | |
| print(f"✅ SAM loaded successfully on {device}") | |
| return global_state["sam_predictor"] | |
| # ------------------------------------------------------------ | |
| # 🖼️ Step 1: Upload Image | |
| # ------------------------------------------------------------ | |
| def process_upload(image): | |
| if image is None: | |
| return None, "⚠️ Upload gambar terlebih dahulu!", None, None | |
| image_np = np.array(image) | |
| global_state["original_image"] = image_np.copy() | |
| predictor = load_sam_model() | |
| predictor.set_image(image_np) | |
| global_state["image_set"] = True | |
| status = "✅ Gambar berhasil diupload!\n\nKlik 'Auto Smart Select' atau gunakan Box/Point untuk memilih objek." | |
| return Image.fromarray(image_np), status, None, None | |
| # ------------------------------------------------------------ | |
| # ✨ Step 2: Auto Smart Select (Pilih Objek Otomatis) | |
| # ------------------------------------------------------------ | |
| def auto_smart_select(): | |
| if not global_state["image_set"]: | |
| return None, None, "⚠️ Upload gambar terlebih dahulu!" | |
| from segment_anything import SamAutomaticMaskGenerator | |
| predictor = global_state["sam_predictor"] | |
| sam_model = predictor.model | |
| mask_generator = SamAutomaticMaskGenerator( | |
| model=sam_model, | |
| points_per_side=32, | |
| pred_iou_thresh=0.88, | |
| stability_score_thresh=0.93, | |
| min_mask_region_area=500 | |
| ) | |
| masks = mask_generator.generate(global_state["original_image"]) | |
| if not masks: | |
| return None, None, "❌ Tidak ada objek terdeteksi." | |
| # Pilih mask terbaik berdasarkan area dan posisi tengah | |
| h, w = global_state["original_image"].shape[:2] | |
| center_x, center_y = w // 2, h // 2 | |
| def score_mask(mask_data): | |
| mask = mask_data['segmentation'] | |
| y, x = np.where(mask) | |
| if len(x) == 0: return 0 | |
| dist = np.sqrt((x.mean() - center_x)**2 + (y.mean() - center_y)**2) | |
| dist_score = 1 - (dist / np.sqrt(center_x**2 + center_y**2)) | |
| area_ratio = mask_data['area'] / (h * w) | |
| size_score = 1 - abs(area_ratio - 0.15) # ideal sekitar 15% | |
| return dist_score * 0.6 + size_score * 0.4 | |
| best_mask = max(masks, key=score_mask)['segmentation'] | |
| rgba_image = np.dstack((global_state["original_image"], (best_mask * 255).astype(np.uint8))) | |
| global_state["current_mask"] = best_mask | |
| global_state["rgba_image"] = rgba_image | |
| segmented_preview = Image.fromarray(rgba_image) | |
| mask_preview = Image.fromarray((best_mask * 255).astype(np.uint8)) | |
| status = "✅ Foreground otomatis terpilih!\nKlik 'Generate Background' untuk mengganti latar belakang." | |
| return segmented_preview, mask_preview, status | |
| # ------------------------------------------------------------ | |
| # 🎨 Step 3: Ganti Background dengan Prompt | |
| # ------------------------------------------------------------ | |
| def change_background(background_prompt, negative_prompt, guidance_scale, num_steps, seed): | |
| if global_state["rgba_image"] is None: | |
| return None, "⚠️ Generate mask dulu!" | |
| rgba_image = Image.fromarray(global_state["rgba_image"]) | |
| pipeline = StableDiffusionXLInpaintPipeline.from_pretrained( | |
| "diffusers/stable-diffusion-xl-1.0-inpainting-0.1", | |
| torch_dtype=torch.float16, | |
| variant="fp16" | |
| ) | |
| if torch.cuda.is_available(): | |
| pipeline = pipeline.to("cuda") | |
| alpha = np.array(rgba_image.split()[-1]) | |
| mask = np.where(alpha < 128, 255, 0).astype(np.uint8) | |
| mask = Image.fromarray(mask).convert("L") | |
| rgb_image = rgba_image.convert("RGB") | |
| generator = torch.manual_seed(int(seed)) | |
| result = pipeline( | |
| image=rgb_image, | |
| mask_image=mask, | |
| prompt=background_prompt, | |
| negative_prompt=negative_prompt, | |
| guidance_scale=guidance_scale, | |
| generator=generator, | |
| num_inference_steps=int(num_steps), | |
| width=rgb_image.width, | |
| height=rgb_image.height | |
| ).images[0] | |
| status = f"✅ Background diganti: {background_prompt}" | |
| return result, status | |
| # ------------------------------------------------------------ | |
| # 🧱 Gradio UI | |
| # ------------------------------------------------------------ | |
| def create_gradio_interface(): | |
| with gr.Blocks(title="AI Background Maker (Offline SAM)") as demo: | |
| gr.Markdown("# 🖼️ AI Background Maker\nUpload foto → Pilih objek → Ganti background ✨") | |
| with gr.Row(): | |
| with gr.Column(scale=1): | |
| input_image = gr.Image(label="Upload Gambar", type="pil") | |
| upload_btn = gr.Button("📥 Load Image", variant="primary") | |
| auto_smart_btn = gr.Button("🤖 Auto Smart Select (No Click!)", variant="primary") | |
| bg_prompt = gr.Textbox(value="sunset beach, cinematic, 8k", label="Background Prompt") | |
| neg_prompt = gr.Textbox(value="blurry, low quality", label="Negative Prompt") | |
| guidance = gr.Slider(1, 15, 7, 0.5, label="Guidance Scale") | |
| steps = gr.Slider(10, 50, 30, 5, label="Steps") | |
| seed = gr.Number(value=42, label="Seed") | |
| gen_bg = gr.Button("🚀 Generate Background", variant="secondary") | |
| with gr.Column(scale=1): | |
| status = gr.Textbox(label="Status", lines=5) | |
| segmented = gr.Image(label="Foreground") | |
| mask = gr.Image(label="Mask") | |
| result = gr.Image(label="Final Output") | |
| result_status = gr.Textbox(label="Output Status", lines=2) | |
| upload_btn.click(process_upload, inputs=[input_image], | |
| outputs=[segmented, status, mask, result]) | |
| auto_smart_btn.click(auto_smart_select, outputs=[segmented, mask, status]) | |
| gen_bg.click(change_background, | |
| inputs=[bg_prompt, neg_prompt, guidance, steps, seed], | |
| outputs=[result, result_status]) | |
| return demo | |
| if __name__ == "__main__": | |
| demo = create_gradio_interface() | |
| demo.launch(share=True, debug=True, server_name="0.0.0.0", server_port=7860) | |