Update app.py

app.py

@@ -71,11 +71,9 @@ except Exception as e:
     print(f"Error loading model: {str(e)}")
     traceback.print_exc()
 
-# Define image transformation pipeline
 transform = transforms.Compose([
     transforms.Resize((128, 128)),
-    transforms.PILToTensor(),  # Changed from ToTensor()
-    transforms.ConvertImageDtype(torch.float32),
+    transforms.ToTensor(),  # PILToTensor 대신 ToTensor 사용
     transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
 ])
 
@@ -86,14 +84,21 @@ def process_image(image):
     try:
         # Convert numpy array to PIL Image
         if isinstance(image, np.ndarray):
-            image = Image.fromarray(image.astype('uint8'))
+            # Ensure the array is uint8
+            if image.dtype != np.uint8:
+                image = (image * 255).astype(np.uint8)
+            image = Image.fromarray(image)
         
         # Convert to RGB if necessary
         if image.mode != 'RGB':
             image = image.convert('RGB')
         
+        # Resize the image
+        image = image.resize((128, 128), Image.Resampling.LANCZOS)
+        
         print(f"Processed image size: {image.size}")
         print(f"Processed image mode: {image.mode}")
+        print(f"Image type: {type(image)}")
         
         return image
     except Exception as e:
@@ -111,10 +116,15 @@ def predict(image):
         if processed_image is None:
             return {cls: 0.0 for cls in ["Rope", "Hammer", "Other"]}
         
-        # Transform image to tensor using torchvision transforms
+        # Transform image to tensor
         try:
-            tensor_image = transform(processed_image).unsqueeze(0)
+            # Convert PIL Image to tensor
+            tensor_image = transform(processed_image)
+            # Add batch dimension
+            tensor_image = tensor_image.unsqueeze(0)
             print(f"Input tensor shape: {tensor_image.shape}")
+            print(f"Tensor dtype: {tensor_image.dtype}")
+            print(f"Tensor device: {tensor_image.device}")
         except Exception as e:
             print(f"Error in tensor conversion: {str(e)}")
             traceback.print_exc()
@@ -142,7 +152,7 @@ def predict(image):
 # Gradio interface
 interface = gr.Interface(
     fn=predict,
-    inputs=gr.Image(),
+    inputs=gr.Image(type="pil"),  # PIL 이미지 타입으로 명시
     outputs=gr.Label(num_top_classes=3),
     title="Mechanical Tools Classifier",
     description="Upload an image of a tool to classify it as 'Rope', 'Hammer', or 'Other'.",