matte

app.py

@@ -6,6 +6,7 @@ from color_match_processor import create_color_match_tab
 from simple_effects_processor import create_effects_tab
 from histogram_processor import create_histogram_tab
 from blend_processor import create_blend_tab
+from matte_processor import create_matte_tab
 
 with gr.Blocks(title="Image Processing Suite") as demo:
     gr.Markdown("# Image Processing Suite")
@@ -17,6 +18,7 @@ with gr.Blocks(title="Image Processing Suite") as demo:
     create_effects_tab()
     create_histogram_tab()
     create_blend_tab()
+    create_matte_tab()  # Add this line
     
-
-demo.launch(debug=True)
+if __name__ == "__main__":
+    demo.launch(share=True)

matte_processor.py

@@ -0,0 +1,120 @@
+import gradio as gr
+from PIL import Image
+import torch
+from transformers import VitMatteImageProcessor, VitMatteForImageMatting
+import math
+from pathlib import Path
+
+class VITMatteModel:
+    def __init__(self, model, processor):
+        self.model = model
+        self.processor = processor
+
+def load_VITMatte_model(local_files_only=False):
+    model = VitMatteForImageMatting.from_pretrained("hustvl/vitmatte-small-composition-1k", local_files_only=local_files_only)
+    processor = VitMatteImageProcessor.from_pretrained("hustvl/vitmatte-small-composition-1k", local_files_only=local_files_only)
+    return VITMatteModel(model, processor)
+
+def generate_VITMatte(image, trimap, local_files_only=False, device="cpu", max_megapixels=2.0):
+    if image is None or trimap is None:
+        return None
+        
+    # Convert to proper formats
+    if isinstance(image, np.ndarray):
+        image = Image.fromarray(image)
+    if isinstance(trimap, np.ndarray):
+        trimap = Image.fromarray(trimap)
+        
+    if image.mode != 'RGB':
+        image = image.convert('RGB')
+    if trimap.mode != 'L':
+        trimap = trimap.convert('L')
+
+    # Calculate resize if needed
+    max_megapixels *= 1048576
+    width, height = image.size
+    ratio = width / height
+    target_width = math.sqrt(ratio * max_megapixels)
+    target_height = target_width / ratio
+    target_width = int(target_width)
+    target_height = int(target_height)
+    
+    resized = False
+    if width * height > max_megapixels:
+        image = image.resize((target_width, target_height), Image.BILINEAR)
+        trimap = trimap.resize((target_width, target_height), Image.BILINEAR)
+        resized = True
+
+    # Set device
+    if device == "cuda" and not torch.cuda.is_available():
+        device = "cpu"
+    device = torch.device(device)
+
+    # Load and process
+    vit_matte_model = load_VITMatte_model(local_files_only=local_files_only)
+    vit_matte_model.model.to(device)
+
+    inputs = vit_matte_model.processor(images=image, trimaps=trimap, return_tensors="pt")
+    
+    with torch.no_grad():
+        inputs = {k: v.to(device) for k, v in inputs.items()}
+        predictions = vit_matte_model.model(**inputs).alphas
+
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+        torch.cuda.ipc_collect()
+
+    # Convert prediction to image
+    mask = predictions.cpu().squeeze().numpy()
+    mask = (mask * 255).astype(np.uint8)
+    mask = Image.fromarray(mask).convert('L')
+    
+    mask = mask.crop((0, 0, image.width, image.height))
+
+    if resized:
+        mask = mask.resize((width, height), Image.BILINEAR)
+
+    return np.array(mask)
+
+def create_matte_tab():
+    with gr.Tab("Image Matting"):
+        with gr.Row():
+            with gr.Column():
+                input_image = gr.Image(label="Input Image", type="numpy", height=256)
+                trimap_image = gr.Image(label="Trimap Image", type="numpy", height=256)
+                
+                device = gr.Radio(
+                    choices=["cpu", "cuda"], 
+                    value="cpu",
+                    label="Device"
+                )
+                
+                max_megapixels = gr.Slider(
+                    minimum=0.5,
+                    maximum=8.0,
+                    value=2.0,
+                    step=0.5,
+                    label="Max Megapixels"
+                )
+                
+                local_files = gr.Checkbox(
+                    value=False,
+                    label="Use Local Files Only"
+                )
+                
+                process_btn = gr.Button("Generate Matte")
+            
+            with gr.Column():
+                output_image = gr.Image(label="Generated Matte")
+        
+        process_btn.click(
+            fn=generate_VITMatte,
+            inputs=[
+                input_image,
+                trimap_image,
+                local_files,
+                device,
+                max_megapixels
+            ],
+            outputs=output_image
+        )