Update app.py

app.py

@@ -71,11 +71,10 @@ except Exception as e:
     print(f"Error loading model: {str(e)}")
     traceback.print_exc()
 
-# Define image transformation pipeline to match training
 transform = transforms.Compose([
     transforms.Resize((128, 128)),
+    transforms.Lambda(lambda x: x.convert('RGB')),
     transforms.ToTensor(),
-    # Using standard normalization as in training
     transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
 ])
 
@@ -84,15 +83,17 @@ def process_image(image):
         return None
     
     try:
-        # Convert numpy array to PIL Image
+        # numpy array를 PIL Image로 변환
         if isinstance(image, np.ndarray):
-            image = Image.fromarray(image)
+            image = Image.fromarray(image.astype('uint8'))
         
-        # Convert to RGB if necessary
+        # 이미지가 RGB가 아니면 변환
         if image.mode != 'RGB':
             image = image.convert('RGB')
-            
-        # Print debug information
+        
+        # 이미지 크기 조정
+        image = image.resize((128, 128), Image.Resampling.LANCZOS)
+        
         print(f"Processed image size: {image.size}")
         print(f"Processed image mode: {image.mode}")
         
@@ -107,16 +108,32 @@ def predict(image):
         return {cls: 0.0 for cls in ["Rope", "Hammer", "Other"]}
     
     try:
-        # Process the image
+        # 이미지 처리
         processed_image = process_image(image)
         if processed_image is None:
             return {cls: 0.0 for cls in ["Rope", "Hammer", "Other"]}
         
-        # Transform for model
-        tensor_image = transform(processed_image).unsqueeze(0)
-        print(f"Input tensor shape: {tensor_image.shape}")
+        # PIL Image를 텐서로 변환
+        try:
+            # PIL Image를 numpy array로 변환
+            img_array = np.array(processed_image)
+            # numpy array를 torch tensor로 변환
+            tensor_image = torch.from_numpy(img_array.transpose((2, 0, 1))).float() / 255.0
+            # 정규화
+            tensor_image = transforms.Normalize(
+                mean=[0.485, 0.456, 0.406],
+                std=[0.229, 0.224, 0.225]
+            )(tensor_image)
+            # 배치 차원 추가
+            tensor_image = tensor_image.unsqueeze(0)
+            
+            print(f"Input tensor shape: {tensor_image.shape}")
+        except Exception as e:
+            print(f"Error in tensor conversion: {str(e)}")
+            traceback.print_exc()
+            return {cls: 0.0 for cls in ["Rope", "Hammer", "Other"]}
         
-        # Make prediction
+        # 예측 수행
         with torch.no_grad():
             outputs = model(tensor_image)
             print(f"Raw outputs: {outputs}")
@@ -124,7 +141,7 @@ def predict(image):
             probabilities = F.softmax(outputs, dim=1)[0].cpu().numpy()
             print(f"Probabilities: {probabilities}")
             
-        # Return results
+        # 결과 반환
         classes = ["Rope", "Hammer", "Other"]
         results = {cls: float(prob) for cls, prob in zip(classes, probabilities)}
         print(f"Final results: {results}")
@@ -135,15 +152,15 @@ def predict(image):
         traceback.print_exc()
         return {cls: 0.0 for cls in ["Rope", "Hammer", "Other"]}
 
-# Gradio interface
+# Gradio 인터페이스
 interface = gr.Interface(
     fn=predict,
-    inputs=gr.Image(),  # Accept any image format
+    inputs=gr.Image(),
     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'.",
 )
 
-# Launch the interface
+# 인터페이스 실행
 if __name__ == "__main__":
     interface.launch()