Create app.py

app.py

@@ -0,0 +1,207 @@
+import numpy as np
+import torch
+import torch.nn.functional as F
+from torchvision.transforms.functional import normalize
+import gradio as gr
+from PIL import Image
+from typing import Tuple, Dict, List
+import cv2
+from pathlib import Path
+
+from advanced_briarmb import MultiTargetBriaRMBG, ClothingType, GarmentFeatures
+
+class ImageProcessor:
+    def __init__(self, model_path: str = "briaai/RMBG-1.4"):
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.net = MultiTargetBriaRMBG.from_pretrained(model_path)
+        self.net.to(self.device)
+        self.model_input_size = (1024, 1024)
+        
+    def preprocess_image(self, image: np.ndarray) -> torch.Tensor:
+        """Prepare image for model input"""
+        # Convert numpy array to PIL
+        if isinstance(image, np.ndarray):
+            image = Image.fromarray(image)
+        
+        # Convert to RGB and resize
+        image = image.convert('RGB')
+        image = image.resize(self.model_input_size, Image.LANCZOS)
+        
+        # Convert to tensor
+        im_tensor = torch.tensor(np.array(image), dtype=torch.float32).permute(2, 0, 1)
+        im_tensor = torch.unsqueeze(im_tensor, 0)
+        im_tensor = torch.divide(im_tensor, 255.0)
+        im_tensor = normalize(im_tensor, [0.5, 0.5, 0.5], [1.0, 1.0, 1.0])
+        
+        return im_tensor.to(self.device)
+
+    def postprocess_mask(self, mask: torch.Tensor, target_size: Tuple[int, int]) -> Image.Image:
+        """Convert model output mask to PIL Image"""
+        # Resize mask to original image size
+        mask = F.interpolate(mask, size=target_size, mode='bilinear')
+        mask = torch.squeeze(mask, 0)
+        
+        # Normalize mask values
+        mask = (mask - mask.min()) / (mask.max() - mask.min())
+        
+        # Convert to PIL Image
+        mask_np = (mask * 255).cpu().data.numpy().astype(np.uint8)
+        return Image.fromarray(mask_np[0])
+
+    def process_image(self, 
+                     image: np.ndarray,
+                     mode: str = "background_removal",
+                     clothing_options: Dict = None) -> Dict[str, Image.Image]:
+        """Main processing function"""
+        # Get original size
+        orig_image = Image.fromarray(image) if isinstance(image, np.ndarray) else image
+        orig_size = orig_image.size
+        
+        # Preprocess
+        input_tensor = self.preprocess_image(image)
+        
+        # Model inference
+        results = self.net(input_tensor, mode=mode)
+        
+        # Process different outputs based on mode
+        outputs = {}
+        
+        if mode == "background_removal" or mode == "all":
+            # Get foreground mask
+            fg_mask = self.postprocess_mask(results["foreground"], orig_size)
+            
+            # Create transparent background image
+            transparent = Image.new("RGBA", orig_size, (0, 0, 0, 0))
+            transparent.paste(orig_image, mask=fg_mask)
+            outputs["removed_background"] = transparent
+            
+            # Extract background if requested
+            if mode == "all":
+                bg_mask = self.postprocess_mask(results["background"], orig_size)
+                background = Image.new("RGBA", orig_size, (0, 0, 0, 0))
+                background.paste(orig_image, mask=bg_mask)
+                outputs["background"] = background
+        
+        if mode == "clothing" or mode == "all":
+            clothing_mask = self.postprocess_mask(results["clothing"], orig_size)
+            
+            if clothing_options:
+                # Apply clothing modifications
+                modified = self.apply_clothing_modifications(
+                    orig_image,
+                    clothing_mask,
+                    clothing_options
+                )
+                outputs["modified_clothing"] = modified
+            
+            # Extract original clothing
+            clothing = Image.new("RGBA", orig_size, (0, 0, 0, 0))
+            clothing.paste(orig_image, mask=clothing_mask)
+            outputs["clothing"] = clothing
+            
+        return outputs
+
+    def apply_clothing_modifications(self,
+                                   image: Image.Image,
+                                   mask: Image.Image,
+                                   options: Dict) -> Image.Image:
+        """Apply clothing modifications based on options"""
+        if "color" in options:
+            image = self.change_clothing_color(image, mask, options["color"])
+        
+        if "pattern" in options:
+            image = self.apply_pattern(image, mask, options["pattern"])
+            
+        if "style_transfer" in options:
+            image = self.transfer_clothing_style(image, mask, options["style_transfer"])
+            
+        return image
+
+def create_ui() -> gr.Blocks:
+    """Create the Gradio UI"""
+    processor = ImageProcessor()
+    
+    with gr.Blocks() as app:
+        gr.Markdown("# Advanced Background and Clothing Removal")
+        
+        with gr.Tab("Background Removal"):
+            with gr.Row():
+                with gr.Column():
+                    input_image = gr.Image(label="Input Image", type="numpy")
+                    remove_bg_btn = gr.Button("Remove Background")
+                
+                with gr.Column():
+                    output_image = gr.Image(label="Result", type="pil")
+            
+            remove_bg_btn.click(
+                fn=lambda img: processor.process_image(img)["removed_background"],
+                inputs=[input_image],
+                outputs=[output_image]
+            )
+        
+        with gr.Tab("Clothing Manipulation"):
+            with gr.Row():
+                with gr.Column():
+                    cloth_input = gr.Image(label="Input Image", type="numpy")
+                    
+                    with gr.Accordion("Clothing Options"):
+                        color_picker = gr.ColorPicker(label="New Color")
+                        pattern_choice = gr.Dropdown(
+                            choices=["Stripes", "Dots", "Floral"],
+                            label="Pattern"
+                        )
+                        style_image = gr.Image(label="Style Reference", type="numpy")
+                    
+                    process_clothing_btn = gr.Button("Process Clothing")
+                
+                with gr.Column():
+                    cloth_output = gr.Image(label="Modified Clothing", type="pil")
+            
+            def process_clothing(image, color, pattern, style):
+                options = {}
+                if color:
+                    options["color"] = color
+                if pattern:
+                    options["pattern"] = pattern
+                if style is not None:
+                    options["style_transfer"] = style
+                    
+                return processor.process_image(
+                    image,
+                    mode="clothing",
+                    clothing_options=options
+                )["modified_clothing"]
+            
+            process_clothing_btn.click(
+                fn=process_clothing,
+                inputs=[cloth_input, color_picker, pattern_choice, style_image],
+                outputs=[cloth_output]
+            )
+        
+        # Examples section
+        examples_dir = Path("./examples")
+        examples = [
+            [str(examples_dir / f"example_{i}.jpg")]
+            for i in range(1, 4)
+            if (examples_dir / f"example_{i}.jpg").exists()
+        ]
+        
+        if examples:
+            gr.Examples(
+                examples=examples,
+                inputs=[input_image],
+                outputs=[output_image],
+                fn=lambda img: processor.process_image(img)["removed_background"],
+                cache_examples=True
+            )
+    
+    return app
+
+if __name__ == "__main__":
+    app = create_ui()
+    app.launch(
+        share=False,
+        server_name="0.0.0.0",
+        server_port=7860,
+        debug=True
+    )