sudip2003 commited on
Commit
cdd2881
·
verified ·
1 Parent(s): 31eeb66

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +39 -39
app.py CHANGED
@@ -1,39 +1,39 @@
1
- import gradio as gr
2
- import cv2
3
- import numpy as np
4
- from ultralytics import YOLO
5
-
6
-
7
- model = YOLO(r"D:\Full_stack_ML\Internship\U\best.pt")
8
-
9
- def process_image(image):
10
- image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
11
- results = model.predict(image, conf=0.15)
12
-
13
- if len(results[0].boxes.cls) == 1:
14
- mask_tensor = results[0].masks.data[0].cpu().numpy()
15
- mask = (mask_tensor * 255).astype(np.uint8)
16
-
17
- mask = cv2.resize(mask, (image.shape[1], image.shape[0]))
18
- kernel = np.ones((5, 5), np.uint8)
19
- mask = cv2.dilate(mask, kernel, iterations=2)
20
- mask = cv2.erode(mask, kernel, iterations=2)
21
-
22
- rgba_image = cv2.cvtColor(image, cv2.COLOR_BGR2RGBA)
23
- rgba_image[:, :, 3] = mask
24
-
25
- return rgba_image
26
- else:
27
- return "Error: Uploaded image has more than one face. Please upload a different image."
28
-
29
-
30
- demo = gr.Interface(
31
- fn=process_image,
32
- inputs=gr.Image(type="numpy"),
33
- outputs=gr.Image(type="numpy"),
34
- title="Face Segmentation",
35
- description="Upload an image"
36
- )
37
-
38
- if __name__ == "__main__":
39
- demo.launch()
 
1
+ import gradio as gr
2
+ import cv2
3
+ import numpy as np
4
+ from ultralytics import YOLO
5
+
6
+
7
+ model = YOLO(r"best.pt")
8
+
9
+ def process_image(image):
10
+ image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
11
+ results = model.predict(image, conf=0.15)
12
+
13
+ if len(results[0].boxes.cls) == 1:
14
+ mask_tensor = results[0].masks.data[0].cpu().numpy()
15
+ mask = (mask_tensor * 255).astype(np.uint8)
16
+
17
+ mask = cv2.resize(mask, (image.shape[1], image.shape[0]))
18
+ kernel = np.ones((5, 5), np.uint8)
19
+ mask = cv2.dilate(mask, kernel, iterations=2)
20
+ mask = cv2.erode(mask, kernel, iterations=2)
21
+
22
+ rgba_image = cv2.cvtColor(image, cv2.COLOR_BGR2RGBA)
23
+ rgba_image[:, :, 3] = mask
24
+
25
+ return rgba_image
26
+ else:
27
+ return "Error: Uploaded image has more than one face. Please upload a different image."
28
+
29
+
30
+ demo = gr.Interface(
31
+ fn=process_image,
32
+ inputs=gr.Image(type="numpy"),
33
+ outputs=gr.Image(type="numpy"),
34
+ title="Face Segmentation",
35
+ description="Upload an image"
36
+ )
37
+
38
+ if __name__ == "__main__":
39
+ demo.launch()