Create app.py

app.py

@@ -0,0 +1,45 @@
+import gradio as gr
+import numpy as np
+import torch
+from transformers import Mask2FormerImageProcessor, Mask2FormerForUniversalSegmentation
+
+# Load the model from Hugging Face
+model_name = "gdurkin/cdl_mask2former_hi_res_v3"
+processor = Mask2FormerImageProcessor.from_pretrained(model_name)
+model = Mask2FormerForUniversalSegmentation.from_pretrained(model_name)
+device = torch.device('cpu')
+
+# Define the inference function
+def predict(img):
+    if isinstance(img, np.ndarray):
+        img = torch.from_numpy(img).float()
+    if torch.is_tensor(img):
+        input_tensor = img
+    else:
+        raise ValueError("Unsupported image format")
+
+    if input_tensor.ndim == 3:
+        input_tensor = input_tensor.unsqueeze(0)
+    elif input_tensor.ndim != 4:
+        raise ValueError("Input tensor must be 3D or 4D")
+    
+    input_tensor = input_tensor.permute(0, 3, 1, 2)  # Ensure the tensor is in the correct shape (N, C, H, W)
+    
+    with torch.no_grad():
+        outputs = model(input_tensor.to(device))
+    
+    target_sizes = [(input_tensor.shape[2], input_tensor.shape[3])]
+    predicted_segmentation_maps = processor.post_process_semantic_segmentation(outputs, target_sizes=target_sizes)
+    
+    return predicted_segmentation_maps[0].cpu().numpy()
+
+# Create a Gradio interface
+iface = gr.Interface(
+    fn=predict,
+    inputs=gr.Image(type="numpy", image_mode='RGB'),
+    outputs="numpy",
+    live=True
+)
+
+# Launch the interface
+iface.launch()