try claude image solution test

handler.py

@@ -1,9 +1,10 @@
-import base64
 import torch
 import torchvision.transforms as transforms
 from PIL import Image
 import io
+import base64
 import logging
+import json
 
 # Set up logging
 logging.basicConfig(level=logging.INFO)
@@ -15,7 +16,6 @@ class EndpointHandler:
         self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
         self.model = torch.jit.load(f"{model_dir}/model_scripted_efficientnet.pt", map_location=self.device)
         self.model.eval()
-
         self.transform = transforms.Compose([
             transforms.Resize((224, 224)),
             transforms.ToTensor(),
@@ -24,7 +24,6 @@ class EndpointHandler:
                 std=[0.229, 0.224, 0.225]
             )
         ])
-
         self.supported_issues = [
             "Dark Spots",
             "Dry Lips",
@@ -40,45 +39,72 @@ class EndpointHandler:
 
     def __call__(self, data):
         logger.info(f"Received data: {type(data)}")
-
         image = None
-
-        if isinstance(data, bytes):
-            logger.info("Input is bytes. Attempting to load image.")
-            image = Image.open(io.BytesIO(data)).convert("RGB")
-        elif isinstance(data, dict) and "inputs" in data:
-            logger.info("Input is a dictionary with 'inputs' key.")
-            if isinstance(data["inputs"], bytes):
-                logger.info("'inputs' value is bytes. Attempting to load image.")
-                image = Image.open(io.BytesIO(data["inputs"])).convert("RGB")
-            elif isinstance(data["inputs"], str):
-                logger.info("'inputs' value is a string. Attempting to decode base64.")
-                try:
-                    image_data = base64.b64decode(data["inputs"])
-                    image = Image.open(io.BytesIO(image_data)).convert("RGB")
-                except Exception as e:
-                    logger.error(f"Failed to decode base64 image: {e}")
-            else:
-                logger.warning("'inputs' value is not bytes or base64 string.")
-        else:
-            logger.error("Unsupported input format.")
-            raise ValueError("Unsupported input format")
+        
+        try:
+            # Handle string input (from Hugging Face interface)
+            if isinstance(data, str):
+                logger.info("Input is string. Attempting to parse as JSON.")
+                data = json.loads(data)
+            
+            # Handle various input formats
+            if isinstance(data, dict):
+                if "inputs" in data:
+                    input_data = data["inputs"]
+                    logger.info(f"Input data type: {type(input_data)}")
+                    
+                    # Handle base64 encoded string
+                    if isinstance(input_data, str):
+                        logger.info("Attempting to decode base64 string")
+                        try:
+                            # Remove potential base64 prefix
+                            if "base64," in input_data:
+                                input_data = input_data.split("base64,")[1]
+                            image_bytes = base64.b64decode(input_data)
+                            image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
+                        except Exception as e:
+                            logger.error(f"Failed to decode base64: {str(e)}")
+                    
+                    # Handle raw bytes
+                    elif isinstance(input_data, bytes):
+                        logger.info("Processing raw bytes input")
+                        image = Image.open(io.BytesIO(input_data)).convert("RGB")
+                    
+                    # Handle list input (from Hugging Face interface)
+                    elif isinstance(input_data, list):
+                        logger.info("Processing list input")
+                        if len(input_data) > 0 and isinstance(input_data[0], str):
+                            try:
+                                # Remove potential base64 prefix
+                                if "base64," in input_data[0]:
+                                    input_data[0] = input_data[0].split("base64,")[1]
+                                image_bytes = base64.b64decode(input_data[0])
+                                image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
+                            except Exception as e:
+                                logger.error(f"Failed to decode base64 from list: {str(e)}")
+            
+            # Handle direct bytes input
+            elif isinstance(data, bytes):
+                logger.info("Processing direct bytes input")
+                image = Image.open(io.BytesIO(data)).convert("RGB")
+                
+        except Exception as e:
+            logger.error(f"Error processing input: {str(e)}")
+            raise ValueError(f"Error processing input: {str(e)}")
 
         if image is None:
-            logger.error("Could not load image from input data.")
-            raise ValueError("Could not load image from input data.")
+            logger.error("Could not load image from input data")
+            raise ValueError("Could not load image from input data")
 
         logger.info("Image loaded successfully. Applying transformations.")
         image_tensor = self.transform(image).unsqueeze(0).to(self.device)
-
+        
         with torch.no_grad():
             logger.info("Running inference.")
             outputs = self.model(image_tensor)
         
-        # Use raw outputs (removing softmax)
         predictions = outputs.squeeze().tolist()
         output = [issue for issue, prob in zip(self.supported_issues, predictions) if prob > 0.5]
-
         logger.info(f"Predictions: {output}")
         return {"predictions": output}