renew

app.py

@@ -236,28 +236,60 @@ class MambaIRShadowRemoval(nn.Module):
     
 # Load the model with weights
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-model = MambaIRShadowRemoval(img_channel=3, width=32, middle_blk_num=1, enc_blk_nums=[1, 1, 1, 1], dec_blk_nums=[1, 1, 1, 1], d_state=64)
-model.load_state_dict(torch.load("shadow_removal_model.pth", map_location=device))
-model.to(device)
-model.eval()
-
-# Define the Gradio function
-transform = transforms.Compose([transforms.ToTensor()])
-
-def remove_shadow(image):
+# Define function to load model with specified weights
+def load_model(weights_path):
+    model = MambaIRShadowRemoval(img_channel=3, width=32, middle_blk_num=1, enc_blk_nums=[1, 1, 1, 1], dec_blk_nums=[1, 1, 1, 1], d_state=64)
+    model.load_state_dict(torch.load(weights_path, map_location=device))
+    model.to(device)
+    model.eval()
+    return model
+
+# Preload models for ISTD+ and SRD
+models = {
+    "ISTD+": load_model("ISTD+.pth"),
+    "SRD": load_model("SRD.pth")
+}
+
+# Define transformation
+transform = transforms.Compose([
+    transforms.ToTensor(),
+])
+
+# Define function to perform shadow removal
+def remove_shadow(image, dataset):
+    model = models[dataset]  # Select the appropriate model based on dataset choice
     input_tensor = transform(image).unsqueeze(0).to(device)
     with torch.no_grad():
         output_tensor = model(input_tensor)
     output_image = transforms.ToPILImage()(output_tensor.squeeze(0).cpu())
     return output_image
 
-# Set up Gradio interface
-iface = gr.Interface(
-    fn=remove_shadow,
-    inputs=gr.Image(type="pil"),
-    outputs=gr.Image(type="pil"),
-    title="Shadow Removal Model",
-    description="Upload an image to remove shadows using the trained model."
-)
+# Define example paths for ISTD+ and SRD
+examples = [
+    ["ISTD+.jpg", "ISTD+"],
+    ["SRD.jpg", "SRD"]
+]
+
+# Gradio Interface with dropdown and examples
+with gr.Blocks() as iface:
+    gr.Markdown("## Shadow Removal Model")
+    gr.Markdown("Upload an image to remove shadows using the trained model. Choose the dataset to load the corresponding weights and example images.")
+    
+    with gr.Row():
+        dataset_choice = gr.Dropdown(["ISTD+", "SRD"], label="Choose Dataset", value="ISTD+")
+    
+    example_image = gr.Image(type="pil", label="Input Image")
+    output_image = gr.Image(type="pil", label="Output Image")
+    
+    # Display examples and map them to dataset and images
+    gr.Examples(
+        examples=examples,
+        inputs=[example_image, dataset_choice],
+        label="Examples",
+    )
+    
+    submit_btn = gr.Button("Submit")
+    
+    submit_btn.click(remove_shadow, inputs=[example_image, dataset_choice], outputs=output_image)
 
 iface.launch()