SOSSY commited on
Commit
1daaec8
·
verified ·
1 Parent(s): e1ea0af

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +0 -17
app.py CHANGED
@@ -3,7 +3,6 @@ from transformers import pipeline
3
  from PIL import Image, ImageFilter
4
  import numpy as np
5
 
6
- # Initialize models with fixed choices
7
  segmentation_model = pipeline("image-segmentation", model="nvidia/segformer-b1-finetuned-cityscapes-1024-1024")
8
  depth_estimator = pipeline("depth-estimation", model="Intel/zoedepth-nyu-kitti")
9
 
@@ -25,47 +24,34 @@ def process_image(input_image, method, blur_intensity):
25
  - output_image: final composited image.
26
  - mask_image: the mask used (binary for segmentation, normalized depth for depth-based).
27
  """
28
- # Ensure image is in RGB mode
29
  input_image = input_image.convert("RGB")
30
 
31
  if method == "Segmentation Blur Model":
32
- # Use segmentation to obtain a foreground mask
33
  results = segmentation_model(input_image)
34
- # Assume the last result is the main foreground object
35
  foreground_mask = results[-1]["mask"]
36
- # Ensure the mask is grayscale
37
  foreground_mask = foreground_mask.convert("L")
38
- # Threshold to create a binary mask
39
  binary_mask = foreground_mask.point(lambda p: 255 if p > 128 else 0)
40
 
41
- # Blur the background using Gaussian blur
42
  blurred_background = input_image.filter(ImageFilter.GaussianBlur(radius=blur_intensity))
43
 
44
- # Composite the final image: keep foreground and use blurred background elsewhere
45
  output_image = Image.composite(input_image, blurred_background, binary_mask)
46
  mask_image = binary_mask
47
 
48
  elif method == "Monocular Depth Estimation Model":
49
- # Generate depth map
50
  depth_results = depth_estimator(input_image)
51
  depth_map = depth_results["depth"]
52
 
53
- # Convert depth map to numpy array and normalize to [0, 255]
54
  depth_array = np.array(depth_map).astype(np.float32)
55
  norm = (depth_array - depth_array.min()) / (depth_array.max() - depth_array.min() + 1e-8)
56
  normalized_depth = (norm * 255).astype(np.uint8)
57
  mask_image = Image.fromarray(normalized_depth)
58
 
59
- # Create fully blurred version using Gaussian blur
60
  blurred_image = input_image.filter(ImageFilter.GaussianBlur(radius=blur_intensity))
61
 
62
- # Convert images to arrays for blending
63
  orig_np = np.array(input_image).astype(np.float32)
64
  blur_np = np.array(blurred_image).astype(np.float32)
65
- # Reshape mask for broadcasting
66
  alpha = normalized_depth[..., np.newaxis] / 255.0
67
 
68
- # Blend pixels: 0 = original; 1 = fully blurred
69
  blended_np = (1 - alpha) * orig_np + alpha * blur_np
70
  blended_np = np.clip(blended_np, 0, 255).astype(np.uint8)
71
  output_image = Image.fromarray(blended_np)
@@ -76,7 +62,6 @@ def process_image(input_image, method, blur_intensity):
76
 
77
  return output_image, mask_image
78
 
79
- # Build a Gradio interface
80
  with gr.Blocks() as demo:
81
  gr.Markdown("## FocusFusion: Segmentation & Depth Blur")
82
 
@@ -93,12 +78,10 @@ with gr.Blocks() as demo:
93
  output_image = gr.Image(label="Output Image")
94
  mask_output = gr.Image(label="Mask")
95
 
96
- # Set up event handler
97
  run_button.click(
98
  fn=process_image,
99
  inputs=[input_image, method, blur_intensity],
100
  outputs=[output_image, mask_output]
101
  )
102
 
103
- # Launch the app
104
  demo.launch()
 
3
  from PIL import Image, ImageFilter
4
  import numpy as np
5
 
 
6
  segmentation_model = pipeline("image-segmentation", model="nvidia/segformer-b1-finetuned-cityscapes-1024-1024")
7
  depth_estimator = pipeline("depth-estimation", model="Intel/zoedepth-nyu-kitti")
8
 
 
24
  - output_image: final composited image.
25
  - mask_image: the mask used (binary for segmentation, normalized depth for depth-based).
26
  """
 
27
  input_image = input_image.convert("RGB")
28
 
29
  if method == "Segmentation Blur Model":
 
30
  results = segmentation_model(input_image)
 
31
  foreground_mask = results[-1]["mask"]
 
32
  foreground_mask = foreground_mask.convert("L")
 
33
  binary_mask = foreground_mask.point(lambda p: 255 if p > 128 else 0)
34
 
 
35
  blurred_background = input_image.filter(ImageFilter.GaussianBlur(radius=blur_intensity))
36
 
 
37
  output_image = Image.composite(input_image, blurred_background, binary_mask)
38
  mask_image = binary_mask
39
 
40
  elif method == "Monocular Depth Estimation Model":
 
41
  depth_results = depth_estimator(input_image)
42
  depth_map = depth_results["depth"]
43
 
 
44
  depth_array = np.array(depth_map).astype(np.float32)
45
  norm = (depth_array - depth_array.min()) / (depth_array.max() - depth_array.min() + 1e-8)
46
  normalized_depth = (norm * 255).astype(np.uint8)
47
  mask_image = Image.fromarray(normalized_depth)
48
 
 
49
  blurred_image = input_image.filter(ImageFilter.GaussianBlur(radius=blur_intensity))
50
 
 
51
  orig_np = np.array(input_image).astype(np.float32)
52
  blur_np = np.array(blurred_image).astype(np.float32)
 
53
  alpha = normalized_depth[..., np.newaxis] / 255.0
54
 
 
55
  blended_np = (1 - alpha) * orig_np + alpha * blur_np
56
  blended_np = np.clip(blended_np, 0, 255).astype(np.uint8)
57
  output_image = Image.fromarray(blended_np)
 
62
 
63
  return output_image, mask_image
64
 
 
65
  with gr.Blocks() as demo:
66
  gr.Markdown("## FocusFusion: Segmentation & Depth Blur")
67
 
 
78
  output_image = gr.Image(label="Output Image")
79
  mask_output = gr.Image(label="Mask")
80
 
 
81
  run_button.click(
82
  fn=process_image,
83
  inputs=[input_image, method, blur_intensity],
84
  outputs=[output_image, mask_output]
85
  )
86
 
 
87
  demo.launch()