Update app.py

app.py

@@ -5,6 +5,7 @@ import cv2
 from lang_sam import LangSAM
 from color_matcher import ColorMatcher
 from color_matcher.normalizer import Normalizer
+import torch
 
 # Load the LangSAM model
 model = LangSAM()  # Use the default model or specify custom checkpoint: LangSAM("<model_type>", "<path/to/checkpoint>")
@@ -26,9 +27,12 @@ def apply_color_matching(source_img_np, ref_img_np):
 def extract_and_color_match_sky(image_pil, reference_image_pil, text_prompt="sky"):
     # Use LangSAM to predict the mask for the sky
     masks, boxes, phrases, logits = model.predict(image_pil, text_prompt)
-    
+
+    # Ensure masks is converted from tensor to NumPy
+    masks_np = masks[0].cpu().numpy()  # Convert the tensor to NumPy array
+
     # Convert the mask to a binary format and create a mask image
-    sky_mask = masks[0].astype(np.uint8) * 255
+    sky_mask = (masks_np > 0).astype(np.uint8) * 255  # Ensure it's a binary mask
     
     # Convert PIL image to numpy array for processing
     img_np = np.array(image_pil)
@@ -71,7 +75,7 @@ def gradio_interface():
     outputs = gr.Image(type="pil", label="Resulting Image")
 
     # Launch Gradio app
-    gr.Interface(fn=process_image, inputs=inputs, outputs=outputs, title="Sky Extraction and Color Matching").launch()
+    gr.Interface(fn=process_image, inputs=inputs, outputs=outputs, title="Sky Extraction and Color Matching").launch(share=True)
 
 # Run the Gradio Interface
 if __name__ == "__main__":