Update app.py

app.py

@@ -71,10 +71,11 @@ 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.Lambda(lambda x: x.convert('RGB')),
-    transforms.ToTensor(),
+    transforms.PILToTensor(),  # Changed from ToTensor()
+    transforms.ConvertImageDtype(torch.float32),
     transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
 ])
 
@@ -83,17 +84,14 @@ def process_image(image):
         return None
     
     try:
-        # numpy array를 PIL Image로 변환
+        # Convert numpy array to PIL Image
         if isinstance(image, np.ndarray):
             image = Image.fromarray(image.astype('uint8'))
         
-        # 이미지가 RGB가 아니면 변환
+        # Convert to RGB if necessary
         if image.mode != 'RGB':
             image = image.convert('RGB')
         
-        # 이미지 크기 조정
-        image = image.resize((128, 128), Image.Resampling.LANCZOS)
-        
         print(f"Processed image size: {image.size}")
         print(f"Processed image mode: {image.mode}")
         
@@ -108,32 +106,21 @@ 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"]}
         
-        # PIL Image를 텐서로 변환
+        # Transform image to tensor using torchvision transforms
         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)
-            
+            tensor_image = transform(processed_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}")
@@ -141,7 +128,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}")
@@ -152,7 +139,7 @@ def predict(image):
         traceback.print_exc()
         return {cls: 0.0 for cls in ["Rope", "Hammer", "Other"]}
 
-# Gradio 인터페이스
+# Gradio interface
 interface = gr.Interface(
     fn=predict,
     inputs=gr.Image(),
@@ -161,6 +148,6 @@ interface = gr.Interface(
     description="Upload an image of a tool to classify it as 'Rope', 'Hammer', or 'Other'.",
 )
 
-# 인터페이스 실행
+# Launch the interface
 if __name__ == "__main__":
     interface.launch()