Update app.py

app.py

@@ -6,6 +6,10 @@ from lang_sam import LangSAM
 from color_matcher import ColorMatcher
 from color_matcher.normalizer import Normalizer
 import torch
+import warnings
+
+# Suppress specific warnings if desired
+warnings.filterwarnings("ignore", category=UserWarning)
 
 # Load the LangSAM model
 model = LangSAM()  # Use the default model or specify custom checkpoint if necessary
@@ -15,7 +19,7 @@ def extract_masks(image_pil, prompts):
     masks_dict = {}
     for prompt in prompts_list:
         masks, boxes, phrases, logits = model.predict(image_pil, prompt)
-        if masks:
+        if masks is not None and len(masks) > 0:
             masks_np = masks[0].cpu().numpy()
             mask = (masks_np > 0).astype(np.uint8) * 255  # Binary mask
             masks_dict[prompt] = mask
@@ -142,13 +146,13 @@ def gradio_interface():
                 initial_image = gr.Image(type="pil", label="Upload Image")
                 prompts = gr.Textbox(lines=1, placeholder="Enter prompts separated by commas (e.g., sky, grass)", label="Prompts")
                 segment_button = gr.Button("Segment Image")
-                segment_dropdown = gr.Dropdown(label="Select Segment", choices=[])
+                segment_dropdown = gr.Dropdown(label="Select Segment", choices=[], allow_custom_value=True)
                 replacement_image = gr.Image(type="pil", label="Replacement Image (optional)")
                 color_ref_image = gr.Image(type="pil", label="Color Reference Image (optional)")
                 apply_replacement = gr.Checkbox(label="Apply Replacement", value=False)
                 apply_color_grading = gr.Checkbox(label="Apply Color Grading", value=False)
                 apply_color_to_full_image = gr.Checkbox(label="Apply Color Correction to Full Image", value=False)
-                blending_amount = gr.Slider(minimum=0, maximum=200, step=1, label="Blending Amount", value=0)
+                blending_amount = gr.Slider(minimum=0, maximum=50, step=1, label="Blending Amount", value=0)
                 apply_button = gr.Button("Apply Changes")
                 undo_button = gr.Button("Undo")
             with gr.Column():
@@ -160,9 +164,9 @@ def gradio_interface():
             if initial_image_pil is not None:
                 image_history = [initial_image_pil]
                 current_image_pil = initial_image_pil
-                return current_image_pil, image_history, initial_image_pil, {}, [], "Image loaded."
+                return current_image_pil, image_history, initial_image_pil, {}, gr.Dropdown.update(choices=[], value=None), "Image loaded."
             else:
-                return None, [], None, {}, [], "No image loaded."
+                return None, [], None, {}, gr.Dropdown.update(choices=[], value=None), "No image loaded."
 
         # When the initial image is uploaded, initialize the image history
         initial_image.upload(fn=initialize_image, inputs=initial_image, outputs=[current_image_pil, image_history, current_image_display, masks_dict, segment_dropdown, status])
@@ -170,10 +174,10 @@ def gradio_interface():
         # Segment button click
         def segment_image_wrapper(current_image_pil, prompts):
             if current_image_pil is None:
-                return "No image uploaded.", {}, []
+                return "No image uploaded.", {}, gr.Dropdown.update(choices=[], value=None)
             masks = extract_masks(current_image_pil, prompts)
             if not masks:
-                return "No masks detected for the given prompts.", {}, []
+                return "No masks detected for the given prompts.", {}, gr.Dropdown.update(choices=[], value=None)
             dropdown_choices = list(masks.keys())
             return "Segmentation completed.", masks, gr.Dropdown.update(choices=dropdown_choices, value=dropdown_choices[0])