Update demo

app.py

@@ -8,8 +8,6 @@ import numpy as np
 import gradio as gr
 from seg import U2NETP
 
-DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-
 # Image processing utilities
 def load_image(path: str):
     """ Loads an image from the specified path and converts it to RGB format. """
@@ -34,7 +32,7 @@ class U2NETP_DocSeg(nn.Module):
         return mask
 
 # Initialize the document segmentation model
-docseg = U2NETP_DocSeg(num_classes=1).to(DEVICE)
+docseg = U2NETP_DocSeg(num_classes=1)
 # Load pretrained weights
 docseg_weight_path = './weights/u2netp_docseg_epoch_225_date_2026-01-02.pth'
 checkpoint = torch.load(docseg_weight_path)
@@ -44,7 +42,7 @@ docseg.eval()
 # Get segmentation mask
 def get_mask(image, confidence=0.5):
     org_shape = image.shape[:2]
-    image_tensor = torch.from_numpy(image).float().permute(2, 0, 1).unsqueeze(0).to(DEVICE)
+    image_tensor = torch.from_numpy(image).float().permute(2, 0, 1).unsqueeze(0)
     image_tensor = F.interpolate(image_tensor, size=(288, 288), mode='bilinear')
     with torch.inference_mode():  # faster than no_grad
         mask = docseg(image_tensor)
@@ -53,8 +51,8 @@ def get_mask(image, confidence=0.5):
     return mask[0, 0]  # keep tensor
 
 def overlay_mask(image, mask):
-    image = torch.from_numpy(image).float().to(DEVICE)
-    red = torch.tensor([1.0, 0, 0], device=DEVICE).view(1, 3, 1, 1)
+    image = torch.from_numpy(image).float()
+    red = torch.tensor([1.0, 0, 0]).view(1, 3, 1, 1)
     mask = mask.unsqueeze(0)  # (1, H, W)
     mask = mask.unsqueeze(0)  # (1, 1, H, W)
     overlay = image.permute(2, 0, 1).unsqueeze(0)
@@ -75,4 +73,6 @@ with gr.Blocks() as demo:
         input_image = gr.Image(label="Input Image", type="numpy")
         output_image = gr.Image(label="Segmentation Overlay", type="numpy")
     
-    input_image.change(segment_image, inputs=input_image, outputs=output_image)
+    input_image.change(segment_image, inputs=input_image, outputs=output_image)
+
+demo.launch()