Update app.py

app.py

@@ -2,35 +2,50 @@ import subprocess
 import sys
 import os
 
-# Function to install or reinstall specific packages
-def install(package):
-    subprocess.check_call([sys.executable, "-m", "pip", "install", "--force-reinstall", package])
+# # Function to install or reinstall specific packages
+# def install(package):
+#     subprocess.check_call([sys.executable, "-m", "pip", "install", "--force-reinstall", package])
+
+# # First, ensure NumPy is installed with the correct version
+# try:
+#     import numpy as np
+#     if not np.__version__.startswith("1.24"):
+#         print("Installing compatible NumPy version...")
+#         install("numpy==1.24.3")
+# except ImportError:
+#     print("NumPy not found. Installing...")
+#     install("numpy==1.24.3")
+
+# # Then install other dependencies
+# packages = {
+#     "torch": "2.0.1",
+#     "torchvision": "0.15.2",
+#     "Pillow": "9.5.0",
+#     "gradio": "3.50.2"
+# }
+
+# for package, version in packages.items():
+#     try:
+#         __import__(package.lower())
+#     except ImportError:
+#         print(f"Installing {package}...")
+#         install(f"{package}=={version}")
 
-# First, ensure NumPy is installed with the correct version
-try:
-    import numpy as np
-    if not np.__version__.startswith("1.24"):
-        print("Installing compatible NumPy version...")
-        install("numpy==1.24.3")
-except ImportError:
-    print("NumPy not found. Installing...")
-    install("numpy==1.24.3")
-
-# Then install other dependencies
-packages = {
-    "torch": "2.0.1",
-    "torchvision": "0.15.2",
-    "Pillow": "9.5.0",
-    "gradio": "3.50.2"
-}
-
-for package, version in packages.items():
-    try:
-        __import__(package.lower())
-    except ImportError:
-        print(f"Installing {package}...")
-        install(f"{package}=={version}")
 
+# 먼저 필요한 패키지들을 순서대로 설치
+def install_requirements():
+    packages = [
+        "numpy==1.24.3",
+        "torch==2.0.1",
+        "torchvision==0.15.2",
+        "Pillow==9.5.0",
+        "gradio==3.50.2"
+    ]
+    for package in packages:
+        subprocess.check_call([sys.executable, "-m", "pip", "install", "--force-reinstall", package])
+
+# 패키지 설치 실행
+install_requirements()
 
 import traceback
 import numpy as np
@@ -77,23 +92,39 @@ transform = transforms.Compose([
     transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
 ])
 
+def custom_transform(pil_image):
+    # PIL Image를 numpy array로 변환
+    np_image = np.array(pil_image)
+    
+    # numpy array를 torch tensor로 변환 (채널 순서 변경 포함)
+    tensor_image = torch.from_numpy(np_image.transpose((2, 0, 1))).float()
+    
+    # 값 범위를 [0, 1]로 정규화
+    tensor_image = tensor_image / 255.0
+    
+    # ImageNet 정규화 적용
+    normalize = transforms.Normalize(
+        mean=[0.485, 0.456, 0.406],
+        std=[0.229, 0.224, 0.225]
+    )
+    tensor_image = normalize(tensor_image)
+    
+    return tensor_image
+
 def process_image(image):
     if image is None:
         return None
     
     try:
-        # Convert numpy array to PIL Image
+        # numpy array를 PIL Image로 변환
         if isinstance(image, np.ndarray):
-            # Ensure the array is uint8
-            if image.dtype != np.uint8:
-                image = (image * 255).astype(np.uint8)
-            image = Image.fromarray(image)
+            image = Image.fromarray(image.astype('uint8'))
         
-        # Convert to RGB if necessary
+        # RGB로 변환
         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}")
@@ -111,39 +142,45 @@ 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 image to tensor
         try:
-            # Convert PIL Image to tensor
-            tensor_image = transform(processed_image)
-            # Add batch dimension
-            tensor_image = tensor_image.unsqueeze(0)
+            # 커스텀 변환 함수를 사용하여 텐서로 변환
+            tensor_image = custom_transform(processed_image)
+            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()
             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}")
+        # 예측 수행
+        try:
+            with torch.no_grad():
+                outputs = model(tensor_image)
+                print(f"Raw outputs: {outputs}")
+                
+                probabilities = F.softmax(outputs, dim=1)[0].cpu().numpy()
+                print(f"Probabilities: {probabilities}")
             
-            probabilities = F.softmax(outputs, dim=1)[0].cpu().numpy()
-            print(f"Probabilities: {probabilities}")
+            # 결과 반환
+            classes = ["Rope", "Hammer", "Other"]
+            results = {cls: float(prob) for cls, prob in zip(classes, probabilities)}
+            print(f"Final results: {results}")
+            return results
+            
+        except Exception as e:
+            print(f"Error in prediction: {str(e)}")
+            traceback.print_exc()
+            return {cls: 0.0 for cls in ["Rope", "Hammer", "Other"]}
             
-        # Return results
-        classes = ["Rope", "Hammer", "Other"]
-        results = {cls: float(prob) for cls, prob in zip(classes, probabilities)}
-        print(f"Final results: {results}")
-        return results
-    
     except Exception as e:
         print(f"Prediction error: {str(e)}")
         traceback.print_exc()