Update app.py

app.py

@@ -92,56 +92,48 @@ def apply_depth_based_blur_background(image, mask, strength):
 
 
 def segment_and_blur(input_image, blur_type, gaussian_radius=15, lens_strength=5):
-    """Segments the input image and applies the selected blur."""
-    if oneformer_processor is None or oneformer_model is None:
-        return "Error: OneFormer model not loaded."
+    try:
+        if oneformer_processor is None or oneformer_model is None:
+            return "Error: OneFormer model not loaded."
 
-    image = input_image.convert("RGB")
-    # Rotate the image (assuming this is still needed)
-    image = image.rotate(-90, expand=True)
+        image = input_image.convert("RGB")
+        image = image.rotate(-90, expand=True)
 
-    # Prepare input for semantic segmentation
-    inputs = oneformer_processor(images=image, task_inputs=["semantic"], return_tensors="pt")
+        inputs = oneformer_processor(images=image, task_inputs=["semantic"], return_tensors="pt")
+        with torch.no_grad():
+            outputs = oneformer_model(**inputs)
 
-    # Semantic segmentation
-    with torch.no_grad():
-        outputs = oneformer_model(**inputs)
-
-    # Processing semantic segmentation output
-    predicted_semantic_map = oneformer_processor.post_process_semantic_segmentation(outputs, target_sizes=[image.size[::-1]])[0]
-    segmentation_mask = predicted_semantic_map.cpu().numpy()
+        predicted_semantic_map = oneformer_processor.post_process_semantic_segmentation(outputs, target_sizes=[image.size[::-1]])[0]
+        segmentation_mask = predicted_semantic_map.cpu().numpy()
 
-    # Get the mapping of class IDs to labels
-    id2label = oneformer_model.config.id2label
+        id2label = oneformer_model.config.id2label
+        print(id2label)  # <-- Add this to debug
 
-    # Set foreground label to person
-    foreground_label = 'person'
-    foreground_class_id = None
-    for id, label in id2label.items():
-        if label == foreground_label:
-            foreground_class_id = id
-            break
+        foreground_label = 'person'
+        foreground_class_id = None
+        for id, label in id2label.items():
+            if label.lower() == foreground_label.lower():
+                foreground_class_id = id
+                break
 
-    if foreground_class_id is None:
-        return f"Error: Could not find the label '{foreground_label}' in the model's class mapping."
+        if foreground_class_id is None:
+            return f"Error: Could not find the label '{foreground_label}' in the model's class mapping."
 
-    # Black background mask
-    output_mask_array = np.zeros(segmentation_mask.shape, dtype=np.uint8)
+        output_mask_array = np.zeros(segmentation_mask.shape, dtype=np.uint8)
+        output_mask_array[segmentation_mask == foreground_class_id] = 255
+        mask_pil = Image.fromarray(output_mask_array, mode='L')
 
-    # Set the pixels corresponding to the foreground object to white (255)
-    output_mask_array[segmentation_mask == foreground_class_id] = 255
-
-    # Convert the NumPy array to a PIL Image
-    mask_pil = Image.fromarray(output_mask_array, mode='L')
+        if blur_type == "Gaussian":
+            blurred_image = apply_gaussian_blur_background(image, mask_pil, gaussian_radius)
+        elif blur_type == "Lens":
+            blurred_image = apply_depth_based_blur_background(image, mask_pil, lens_strength)
+        else:
+            return "Error: Invalid blur type selected."
 
-    if blur_type == "Gaussian":
-        blurred_image = apply_gaussian_blur_background(image, mask_pil, gaussian_radius)
-    elif blur_type == "Lens":
-        blurred_image = apply_depth_based_blur_background(image, mask_pil, lens_strength)
-    else:
-        return "Error: Invalid blur type selected."
+        return blurred_image
+    except Exception as e:
+        return f"Error during processing: {str(e)}"
 
-    return blurred_image
 
 iface = gr.Interface(
     fn=segment_and_blur,