Update app.py

app.py

@@ -9,7 +9,78 @@ import io
 import zipfile
 
 # Assuming you have these functions defined elsewhere
-from your_module import preprocess, best_model, DEVICE
+import torch
+import numpy as np
+from PIL import Image
+import albumentations as albu
+import segmentation_models_pytorch as smp
+from albumentations.pytorch.transforms import ToTensorV2
+
+
+
+DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+ENCODER = 'se_resnext50_32x4d'
+ENCODER_WEIGHTS = 'imagenet'
+
+# Load and prepare the model
+best_model = torch.load('deeplabv3+ v15.pth', map_location=DEVICE)
+best_model.eval().float()
+
+def to_tensor(x, **kwargs):
+    return x.astype('float32')#.transpose(2, 0, 1)
+
+# Preprocessing
+preprocessing_fn = smp.encoders.get_preprocessing_fn(ENCODER, ENCODER_WEIGHTS)
+
+def get_preprocessing():
+    _transform = [
+        albu.Resize(512, 512),
+        albu.Lambda(image=preprocessing_fn),
+        albu.Lambda(image=to_tensor, mask=to_tensor),
+        ToTensorV2(),
+        #albu.Normalize(mean=MEAN,std=STD)
+        
+    ]
+    return albu.Compose(_transform)
+
+
+preprocess = get_preprocessing()
+
+@torch.no_grad()
+def process_and_predict(image, model):
+    # Convert PIL Image to numpy array if necessary
+    if isinstance(image, Image.Image):
+        image = np.array(image)
+    
+    # Ensure image is 3-channel
+    if image.ndim == 2:
+        image = np.stack([image] * 3, axis=-1)
+    elif image.shape[2] == 4:
+        image = image[:, :, :3]
+    
+    # Apply preprocessing
+    preprocessed = preprocess(image=image)['image']
+    #preprocessed=torch.tensor(preprocessed)
+    # Add batch dimension and move to device
+    input_tensor = preprocessed.unsqueeze(0).to(DEVICE)
+
+    print(input_tensor.shape)
+    # Predict
+    mask = model(input_tensor)
+    mask = torch.sigmoid(mask)
+    mask = (mask > 0.6).float()
+    
+    # Convert to PIL Image
+    mask_image = Image.fromarray((mask.squeeze().cpu().numpy() * 255).astype(np.uint8))
+    
+    return mask_image
+
+#example
+def main(image_path):
+    image = Image.open(image_path)
+    mask = process_and_predict(image, best_model)
+    return mask
+
 
 def extract_tiles(map_file, model, tile_size=512, overlap=0, batch_size=4):
     tiles = []