Update app.py

app.py

@@ -13,33 +13,40 @@ segmentation_model = BeitForSemanticSegmentation.from_pretrained("microsoft/beit
 depth_feature_extractor = DPTImageProcessor.from_pretrained("Intel/dpt-large")
 depth_model = DPTForDepthEstimation.from_pretrained("Intel/dpt-large")
 
-def apply_gaussian_blur(image):
+def apply_gaussian_blur(image: Image.Image) -> Image.Image:
     # Resize and preprocess the image
     image = image.resize((512, 512)).convert("RGB")
-    inputs = segmentation_processor(image, return_tensors="pt")
+    inputs = image_processor(image, return_tensors="pt")
 
-    # Perform segmentation to get the 'person' mask
+    # Perform semantic segmentation using the model
     with torch.no_grad():
-        outputs = segmentation_model(**inputs)
+        outputs = model(**inputs)
         logits = outputs.logits
-    segmentation = torch.argmax(logits, dim=1)[0]
-    
-    # Create a binary mask for the 'person' class (index 12)
-    person_index = 12
-    binary_mask = (segmentation == person_index).numpy().astype(np.uint8) * 255
 
-    # Apply Gaussian blur to the entire image
+    # Get the predicted class for each pixel
+    segmentation = torch.argmax(logits, dim=1)[0]  # Shape: [height, width]
+
+    # Create a binary mask for the 'person' class
+    person_index = 12  # Assuming 12 is the 'person' class index
+    binary_mask = (segmentation == person_index).numpy().astype(np.uint8) * 255  # Convert to 0 and 255
+
+    # Convert the original image to a numpy array
     image_np = np.array(image)
+
+    # Apply Gaussian blur to the entire image
     blurred_image = cv2.GaussianBlur(image_np, (0, 0), sigmaX=15, sigmaY=15)
 
     # Normalize the mask to range between 0 and 1
     normalized_mask = binary_mask / 255.0
-    normalized_mask = np.expand_dims(normalized_mask, axis=-1)
+    normalized_mask = np.expand_dims(normalized_mask, axis=-1)  # Add channel dimension
 
-    # Create the composite image
+    # Create the composite image with the blurred background
     final_image = (image_np * normalized_mask + blurred_image * (1 - normalized_mask)).astype(np.uint8)
 
-    return Image.fromarray(final_image)
+    # Convert back to PIL Image
+    final_image_pil = Image.fromarray(final_image)
+
+    return final_image_pil
 
 def apply_lens_blur(image):
     # Resize and preprocess the image