Upload handler.py

handler.py

@@ -0,0 +1,115 @@
+# handler.py
+
+import torch
+from transformers import Mask2FormerForUniversalSegmentation, Mask2FormerImageProcessor
+from PIL import Image
+import base64
+import io
+import os
+import numpy as np
+
+class EndpointHandler():
+    def __init__(self, path=""):
+        # Set device
+        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+        # Define label mappings (ensure these match your local environment)
+        self.id2label = {
+            0: 'background',
+            1: 'water',
+            2: 'developed',
+            3: 'corn',
+            4: 'soybeans',
+            5: 'wheat',
+            6: 'other agriculture',
+            7: 'forest/wetlands',
+            8: 'open lands',
+            9: 'barren'
+        }
+        self.label2id = {v: k for k, v in self.id2label.items()}
+
+        # Get the token from environment variables
+        token = os.getenv("HF_API_TOKEN")
+
+        # Load the model with authentication and consistent configurations
+        model_name = "gdurkin/cdl_mask2former_v4_mspc"
+
+        # Initialize the processor and model using from_pretrained
+        self.processor = Mask2FormerImageProcessor.from_pretrained(
+            model_name,
+            use_auth_token=token
+        )
+        self.model = Mask2FormerForUniversalSegmentation.from_pretrained(
+            model_name,
+            use_auth_token=token,
+            id2label=self.id2label,
+            label2id=self.label2id,
+            num_labels=len(self.id2label),
+            ignore_mismatched_sizes=True,
+        )
+        self.model.to(self.device)
+        self.model.eval()
+
+        # Debugging: Print model configuration
+        print("Model configuration:", self.model.config)
+
+    def __call__(self, data):
+        try:
+            # Parse input data
+            if "inputs" in data:
+                image_base64 = data["inputs"]
+            else:
+                return {"error": "No 'inputs' field in request."}
+
+            # Decode the base64 image
+            image_bytes = base64.b64decode(image_base64)
+            image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
+
+            # Convert image to NumPy array and normalize to [0, 1]
+            image_np = np.array(image).astype(np.float32) / 255.0  # Shape: (H, W, C)
+
+            # Convert image to tensor
+            input_tensor = torch.from_numpy(image_np)  # Shape: (H, W, C)
+
+            # Add batch dimension if necessary
+            if input_tensor.ndim == 3:
+                input_tensor = input_tensor.unsqueeze(0)  # Shape: (1, H, W, C)
+            elif input_tensor.ndim != 4:
+                return {"error": "Input tensor must be 3D or 4D"}
+
+            # Permute dimensions to (N, C, H, W)
+            input_tensor = input_tensor.permute(0, 3, 1, 2)
+
+            input_tensor = input_tensor.to(self.device)
+
+            # Perform inference
+            with torch.no_grad():
+                outputs = self.model(pixel_values=input_tensor)
+
+            # Post-process the segmentation map
+            target_sizes = [(input_tensor.shape[2], input_tensor.shape[3])]
+            predicted_segmentation_maps = self.processor.post_process_semantic_segmentation(
+                outputs, target_sizes=target_sizes
+            )
+
+            predicted_segmentation_map = predicted_segmentation_maps[0]  # This is a tensor
+
+
+            # Convert the segmentation map to a NumPy array
+            seg_map_np = predicted_segmentation_map.cpu().numpy()
+         
+            #print("class frequencies:", np.unique(seg_map_np, return_counts=True))
+
+            # Convert the segmentation map to a PNG image
+            seg_map_pil = Image.fromarray(seg_map_np.astype(np.uint8))
+
+            buffered = io.BytesIO()
+            seg_map_pil.save(buffered, format="PNG")
+            seg_map_base64 = base64.b64encode(buffered.getvalue()).decode('utf-8')
+
+            # Return the segmentation map as a base64 string
+            return {'outputs': seg_map_base64}
+
+        except Exception as e:
+            # Handle exceptions and return error message
+            return {"error": str(e)}