Update app.py

app.py

@@ -7,6 +7,8 @@ from PIL import Image
 import os 
 from functools import partial
 
+device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
+print(device)
 
 def resize_image(image, target_size=1024):
     h_img, w_img = image.size
@@ -51,6 +53,7 @@ def segment_image(image, preprocessor, model, crop_size = (1024, 1024), num_clas
           print(x1, x2, y1, y2)
           with torch.no_grad():
             inputs = preprocessor(crop_img, return_tensors = "pt")
+            inputs["pixel_values"] = inputs["pixel_values"].to(device)
             outputs = model(**inputs)
         
           resized_logits = F.interpolate(
@@ -58,7 +61,7 @@ def segment_image(image, preprocessor, model, crop_size = (1024, 1024), num_clas
           )
           preds += F.pad(resized_logits,
                           (int(x1), int(preds.shape[3] - x2), int(y1),
-                          int(preds.shape[2] - y2)))
+                          int(preds.shape[2] - y2))).cpu()
           count_mat[:, :, y1:y2, x1:x2] += 1
     
     assert (count_mat == 0).sum() == 0
@@ -79,8 +82,8 @@ if __name__ == "__main__":
     print(device)
     
     # Load model and processor
-    preprocessor = SegformerImageProcessor.from_pretrained("EPFL-ECEO/segformer-b2-finetuned-coralscapes-1024-1024")
-    model = SegformerForSemanticSegmentation.from_pretrained("EPFL-ECEO/segformer-b2-finetuned-coralscapes-1024-1024").to(device)
+    preprocessor = SegformerImageProcessor.from_pretrained("EPFL-ECEO/segformer-b5-finetuned-coralscapes-1024-1024")
+    model = SegformerForSemanticSegmentation.from_pretrained("EPFL-ECEO/segformer-b5-finetuned-coralscapes-1024-1024").to(device)
     model.eval()
  
     id2label = {"1": "seagrass", "2": "trash", "3": "other coral dead", "4": "other coral bleached", "5": "sand", "6": "other coral alive", "7": "human", "8": "transect tools", "9": "fish", "10": "algae covered substrate", "11": "other animal", "12": "unknown hard substrate", "13": "background", "14": "dark", "15": "transect line", "16": "massive/meandering bleached", "17": "massive/meandering alive", "18": "rubble", "19": "branching bleached", "20": "branching dead", "21": "millepora", "22": "branching alive", "23": "massive/meandering dead", "24": "clam", "25": "acropora alive", "26": "sea cucumber", "27": "turbinaria", "28": "table acropora alive", "29": "sponge", "30": "anemone", "31": "pocillopora alive", "32": "table acropora dead", "33": "meandering bleached", "34": "stylophora alive", "35": "sea urchin", "36": "meandering alive", "37": "meandering dead", "38": "crown of thorn", "39": "dead clam"}