Upload 2 files

app.py

@@ -0,0 +1,466 @@
+import gradio as gr
+import torch
+from PIL import Image
+from torchvision import transforms
+from transformers import AutoModelForImageSegmentation
+import json
+import logging
+import os
+from pathlib import Path
+import io
+import base64
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+# ========== KONFIGURASI ==========
+MODEL_ID = "mohantesting/remove_background"
+MODEL_PATH = "./models/remove_background"
+# ==================================
+
+model = None
+device = None
+transform_image = None
+
+def check_model_exists(path):
+    """Cek apakah model sudah ada"""
+    if not os.path.exists(path):
+        return False
+    
+    required_files = ["config.json"]
+    for file in required_files:
+        if not os.path.exists(os.path.join(path, file)):
+            return False
+    
+    has_weights = False
+    for root, dirs, files in os.walk(path):
+        for file in files:
+            if file.endswith((".bin", ".safetensors")):
+                has_weights = True
+                break
+        if has_weights:
+            break
+    
+    return has_weights
+
+def get_folder_size(folder_path):
+    """Hitung total ukuran folder"""
+    total_size = 0
+    for dirpath, dirnames, filenames in os.walk(folder_path):
+        for filename in filenames:
+            filepath = os.path.join(dirpath, filename)
+            if os.path.isfile(filepath):
+                total_size += os.path.getsize(filepath)
+    return total_size
+
+def download_model():
+    """Download model jika belum ada"""
+    logger.info("="*60)
+    logger.info("CHECKING BACKGROUND REMOVAL MODEL...")
+    logger.info("="*60)
+    
+    if check_model_exists(MODEL_PATH):
+        logger.info("✓ Model sudah ada di local!")
+        logger.info(f"✓ Location: {MODEL_PATH}")
+        
+        size_bytes = get_folder_size(MODEL_PATH)
+        size_mb = size_bytes / (1024 * 1024)
+        logger.info(f"✓ Size: {size_mb:.2f} MB")
+        logger.info("✓ Skipping download...\n")
+        return True
+    
+    logger.info("✗ Model tidak ditemukan. Mulai download...")
+    logger.info(f"Model ID: {MODEL_ID}")
+    logger.info(f"Save to: {MODEL_PATH}")
+    logger.info("-" * 60)
+    
+    try:
+        os.makedirs(MODEL_PATH, exist_ok=True)
+        
+        logger.info("Downloading background removal model...")
+        model_download = AutoModelForImageSegmentation.from_pretrained(
+            MODEL_ID,
+            trust_remote_code=True
+        )
+        model_download.save_pretrained(MODEL_PATH)
+        
+        logger.info("✓ Model downloaded\n")
+        
+        size_bytes = get_folder_size(MODEL_PATH)
+        size_mb = size_bytes / (1024 * 1024)
+        logger.info(f"✓ Total size: {size_mb:.2f} MB")
+        logger.info(f"✓ Model saved at: {MODEL_PATH}\n")
+        
+        del model_download
+        torch.cuda.empty_cache() if torch.cuda.is_available() else None
+        
+        return True
+        
+    except Exception as e:
+        logger.error(f"✗ Error downloading model: {str(e)}")
+        import traceback
+        traceback.print_exc()
+        return False
+
+def load_model():
+    """Load model ke memory"""
+    global model, device, transform_image
+    
+    logger.info("="*60)
+    logger.info("LOADING MODEL INTO MEMORY...")
+    logger.info("="*60)
+    
+    try:
+        device = 'cuda' if torch.cuda.is_available() else 'cpu'
+        logger.info(f"Device: {device}")
+        
+        logger.info("Loading model from local...")
+        model = AutoModelForImageSegmentation.from_pretrained(
+            MODEL_PATH,
+            trust_remote_code=True
+        ).eval().to(device)
+        
+        # Setup transform
+        image_size = (1024, 1024)
+        transform_image = transforms.Compose([
+            transforms.Resize(image_size),
+            transforms.ToTensor(),
+            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
+        ])
+        
+        logger.info("="*60)
+        logger.info("✓ MODEL READY!")
+        logger.info(f"  Model: {MODEL_ID}")
+        logger.info(f"  Device: {device}")
+        logger.info(f"  Image Size: {image_size}")
+        logger.info("="*60 + "\n")
+        
+        return True
+        
+    except Exception as e:
+        logger.error(f"✗ Failed to load model: {str(e)}")
+        import traceback
+        traceback.print_exc()
+        return False
+
+# ========== STARTUP SEQUENCE ==========
+logger.info("Starting application...")
+
+if not download_model():
+    raise Exception("Failed to download model")
+
+if not load_model():
+    raise Exception("Failed to load model into memory")
+
+# ========================================
+
+def remove_background(input_image):
+    """Remove background dari image"""
+    
+    try:
+        if model is None or transform_image is None:
+            return None, None, json.dumps({
+                "success": False,
+                "error": "Model belum siap"
+            }, indent=2, ensure_ascii=False)
+        
+        if input_image is None:
+            return None, None, json.dumps({
+                "success": False,
+                "error": "Image tidak boleh kosong"
+            }, indent=2, ensure_ascii=False)
+        
+        # Convert to PIL Image
+        if not isinstance(input_image, Image.Image):
+            input_image = Image.fromarray(input_image).convert("RGB")
+        else:
+            input_image = input_image.convert("RGB")
+        
+        logger.info(f"Processing image... Size: {input_image.width}x{input_image.height}")
+        
+        # Transform image
+        input_tensor = transform_image(input_image).unsqueeze(0).to(device)
+        
+        # Prediction
+        with torch.no_grad():
+            preds = model(input_tensor)[-1].sigmoid().cpu()
+        
+        pred = preds[0].squeeze()
+        pred_pil = transforms.ToPILImage()(pred)
+        mask = pred_pil.resize(input_image.size)
+        
+        # Create output with alpha channel
+        output_image = input_image.copy()
+        output_image.putalpha(mask)
+        
+        logger.info(f"✓ Background removed. Output: {output_image.width}x{output_image.height}")
+        
+        # JSON result
+        result = {
+            "success": True,
+            "input_size": f"{input_image.width}x{input_image.height}",
+            "output_size": f"{output_image.width}x{output_image.height}",
+            "output_format": "PNG with alpha channel",
+            "model": MODEL_ID,
+            "device": device
+        }
+        
+        return output_image, mask, json.dumps(result, indent=2, ensure_ascii=False)
+    
+    except Exception as e:
+        logger.error(f"Error removing background: {str(e)}", exc_info=True)
+        return None, None, json.dumps({
+            "success": False,
+            "error": str(e)
+        }, indent=2, ensure_ascii=False)
+
+def get_model_info():
+    """Return model info sebagai JSON"""
+    try:
+        info = {
+            "model_name": "Background Removal Model",
+            "model_id": MODEL_ID,
+            "model_path": MODEL_PATH,
+            "model_type": "Image Segmentation",
+            "device": device if device else "unknown",
+            "model_loaded": model is not None,
+            "image_processing_size": "1024x1024",
+            "output_format": "PNG with transparency (alpha channel)",
+            "capabilities": [
+                "Automatic background removal",
+                "High-quality segmentation",
+                "Preserve original image resolution",
+                "Generate alpha mask"
+            ],
+            "use_cases": [
+                "Product photography",
+                "Portrait editing",
+                "E-commerce images",
+                "Graphic design",
+                "Social media content"
+            ]
+        }
+        
+        return json.dumps(info, indent=2, ensure_ascii=False)
+    except Exception as e:
+        return json.dumps({"error": str(e)}, indent=2, ensure_ascii=False)
+
+# Custom CSS
+custom_css = """
+#output_json {
+    font-family: 'Courier New', monospace;
+    font-size: 14px;
+}
+.gradio-container {
+    max-width: 1600px !important;
+}
+"""
+
+# Gradio Interface
+with gr.Blocks(css=custom_css, theme=gr.themes.Soft()) as demo:
+    gr.Markdown("""
+    # 🎨 Background Removal API
+    AI-powered automatic background removal using image segmentation
+    """)
+    
+    with gr.Tabs():
+        # Tab Background Removal
+        with gr.Tab("✂️ Remove Background"):
+            with gr.Row():
+                with gr.Column(scale=1):
+                    input_image = gr.Image(
+                        label="📸 Input Image",
+                        type="pil",
+                        height=450
+                    )
+                    
+                    remove_btn = gr.Button(
+                        "✂️ Remove Background",
+                        variant="primary",
+                        size="lg"
+                    )
+                
+                with gr.Column(scale=1):
+                    output_image = gr.Image(
+                        label="🖼️ Output (No Background)",
+                        type="pil",
+                        height=450
+                    )
+                    
+                    output_mask = gr.Image(
+                        label="🎭 Alpha Mask",
+                        type="pil",
+                        height=200
+                    )
+            
+            output_json = gr.Code(
+                label="📄 JSON Output",
+                language="json",
+                lines=10,
+                elem_id="output_json"
+            )
+            
+            remove_btn.click(
+                fn=remove_background,
+                inputs=[input_image],
+                outputs=[output_image, output_mask, output_json]
+            )
+        
+        # Tab Model Info
+        with gr.Tab("ℹ️ Model Info"):
+            model_info_output = gr.Code(
+                label="Model Information",
+                language="json",
+                lines=30
+            )
+            info_btn = gr.Button("🔍 Get Model Info", variant="secondary")
+            
+            info_btn.click(
+                fn=get_model_info,
+                inputs=[],
+                outputs=model_info_output
+            )
+        
+        # Tab API Documentation
+        with gr.Tab("📚 API Usage"):
+            gr.Markdown("""
+            ## 🚀 API Usage Guide
+            
+            ### 1. Python Example
+            ```python
+            import requests
+            import base64
+            from PIL import Image
+            from io import BytesIO
+            import json
+            
+            # Load and encode image
+            with open("input.jpg", "rb") as f:
+                img_data = base64.b64encode(f.read()).decode()
+            
+            url = "https://YOUR-SPACE-URL/api/predict"
+            
+            payload = {
+                "data": [f"data:image/jpeg;base64,{img_data}"]
+            }
+            
+            response = requests.post(url, json=payload)
+            result = response.json()
+            
+            # Get output image (PNG with transparency)
+            output_image_data = result['data'][0]
+            output_json = json.loads(result['data'][2])
+            
+            # Decode and save
+            img_bytes = base64.b64decode(output_image_data.split(',')[1])
+            img = Image.open(BytesIO(img_bytes))
+            img.save('output_no_bg.png')
+            
+            print(json.dumps(output_json, indent=2))
+            ```
+            
+            ### 2. Response Format
+            ```json
+            {
+              "success": true,
+              "input_size": "1200x1600",
+              "output_size": "1200x1600",
+              "output_format": "PNG with alpha channel",
+              "model": "mohantesting/remove_background",
+              "device": "cuda"
+            }
+            ```
+            
+            ### 3. Output Format
+            
+            - **Format**: PNG with transparency (alpha channel)
+            - **Resolution**: Same as input image
+            - **Background**: Completely transparent
+            - **Quality**: High-quality segmentation
+            
+            ### 4. Best Practices
+            
+            ✅ **DO:**
+            - Use high-resolution images for better results
+            - Ensure good contrast between subject and background
+            - Use well-lit images
+            - Save output as PNG to preserve transparency
+            
+            ❌ **DON'T:**
+            - Don't use extremely large images (>4K) - may cause OOM
+            - Don't expect perfect results on complex backgrounds
+            - Don't save as JPEG (loses transparency)
+            
+            ### 5. Use Cases
+            
+            **E-commerce Product Photos:**
+            ```python
+            # Remove background from product image
+            result = remove_background('product.jpg')
+            result.save('product_no_bg.png')
+            ```
+            
+            **Portrait Photography:**
+            ```python
+            # Remove background from portrait
+            result = remove_background('portrait.jpg')
+            # Can now composite on different backgrounds
+            ```
+            
+            **Graphic Design:**
+            ```python
+            # Create cutouts for design work
+            result = remove_background('subject.jpg')
+            # Use in Photoshop, Canva, etc.
+            ```
+            
+            ### 6. Integration Example
+            
+            **Batch Processing:**
+            ```python
+            import os
+            from pathlib import Path
+            
+            input_dir = 'input_images'
+            output_dir = 'output_images'
+            
+            os.makedirs(output_dir, exist_ok=True)
+            
+            for img_file in Path(input_dir).glob('*.jpg'):
+                result = remove_background(str(img_file))
+                output_path = Path(output_dir) / f"{img_file.stem}_no_bg.png"
+                result.save(output_path)
+                print(f"Processed: {img_file.name}")
+            ```
+            
+            ### 7. Performance
+            
+            - **Processing Time**: ~1-3 seconds per image (GPU)
+            - **Max Resolution**: Recommended up to 2048x2048
+            - **Model Size**: ~180MB
+            - **GPU Memory**: ~2GB recommended
+            
+            ---
+            
+            **Model:** mohantesting/remove_background  
+            **Type:** Image Segmentation  
+            **Framework:** PyTorch + Transformers
+            """)
+    
+    gr.Markdown("""
+    ---
+    ### 💡 Tips for Best Results
+    
+    - Use images with **clear subject-background separation**
+    - **Good lighting** improves accuracy
+    - **Higher resolution** = better edge quality
+    - Save as **PNG** to preserve transparency
+    """)
+
+# Launch
+if __name__ == "__main__":
+    demo.launch(
+        server_name="0.0.0.0",
+        server_port=7860,
+        share=False
+    )

requirements.txt

@@ -0,0 +1,6 @@
+gradio
+torch
+torchvision
+transformers
+Pillow
+accelerate