ihabooe commited on
Commit
0bfc6f3
·
verified ·
1 Parent(s): c6a255a

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +44 -0
app.py ADDED
@@ -0,0 +1,44 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ import torch
3
+ import numpy as np
4
+ from torchvision import transforms
5
+ from PIL import Image
6
+ from torchvision.models.segmentation import deeplabv3_resnet101
7
+
8
+ # Load Pretrained Model
9
+ model = deeplabv3_resnet101(pretrained=True)
10
+ model.eval()
11
+
12
+ def remove_background(image):
13
+ # Preprocess image
14
+ transform = transforms.Compose([
15
+ transforms.ToTensor(),
16
+ transforms.Resize((512, 512)),
17
+ ])
18
+
19
+ input_tensor = transform(image).unsqueeze(0)
20
+
21
+ # Perform inference
22
+ with torch.no_grad():
23
+ output = model(input_tensor)["out"][0]
24
+
25
+ # Create mask (Class 15 = Person)
26
+ mask = output.argmax(0).numpy()
27
+ mask = np.where(mask == 15, 255, 0).astype(np.uint8)
28
+
29
+ # Apply mask
30
+ image = np.array(image)
31
+ transparent_image = np.dstack((image, mask))
32
+
33
+ return Image.fromarray(transparent_image)
34
+
35
+ # Create Gradio Interface
36
+ iface = gr.Interface(
37
+ fn=remove_background,
38
+ inputs=gr.Image(type="pil"),
39
+ outputs=gr.Image(type="pil"),
40
+ title="AI Background Remover",
41
+ description="Upload an image and remove its background using AI."
42
+ )
43
+
44
+ iface.launch()