Update app.py

app.py

@@ -2,37 +2,7 @@ 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])
-
-# # 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",
@@ -44,7 +14,7 @@ def install_requirements():
     for package in packages:
         subprocess.check_call([sys.executable, "-m", "pip", "install", "--force-reinstall", package])
 
-# 패키지 설치 실행
+
 install_requirements()
 
 import traceback
@@ -56,7 +26,7 @@ import torchvision.transforms as transforms
 from PIL import Image
 import gradio as gr
 
-# Define the model exactly as in training
+
 class ModifiedLargeNet(nn.Module):
     def __init__(self):
         super(ModifiedLargeNet, self).__init__()
@@ -75,7 +45,7 @@ class ModifiedLargeNet(nn.Module):
         x = self.fc2(x)
         return x
 
-# Load the trained model with error handling
+
 try:
     model = ModifiedLargeNet()
     state_dict = torch.load("modified_large_net.pt", map_location=torch.device("cpu"))
@@ -88,21 +58,17 @@ except Exception as e:
 
 transform = transforms.Compose([
     transforms.Resize((128, 128)),
-    transforms.ToTensor(),  # PILToTensor 대신 ToTensor 사용
+    transforms.ToTensor(), 
     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]
@@ -116,15 +82,12 @@ def process_image(image):
         return None
     
     try:
-        # numpy array를 PIL Image로 변환
         if isinstance(image, np.ndarray):
             image = Image.fromarray(image.astype('uint8'))
         
-        # RGB로 변환
         if image.mode != 'RGB':
             image = image.convert('RGB')
         
-        # 크기 조정
         image = image.resize((128, 128), Image.Resampling.LANCZOS)
         
         print(f"Processed image size: {image.size}")
@@ -142,15 +105,13 @@ def predict(image):
         return {cls: 0.0 for cls in ["Rope", "Hammer", "Other"]}
     
     try:
-        # 이미지 전처리
         processed_image = process_image(image)
         if processed_image is None:
             return {cls: 0.0 for cls in ["Rope", "Hammer", "Other"]}
         
         try:
-            # 커스텀 변환 함수를 사용하여 텐서로 변환
             tensor_image = custom_transform(processed_image)
-            tensor_image = tensor_image.unsqueeze(0)  # 배치 차원 추가
+            tensor_image = tensor_image.unsqueeze(0) 
             
             print(f"Input tensor shape: {tensor_image.shape}")
             print(f"Tensor dtype: {tensor_image.dtype}")
@@ -161,7 +122,6 @@ def predict(image):
             traceback.print_exc()
             return {cls: 0.0 for cls in ["Rope", "Hammer", "Other"]}
         
-        # 예측 수행
         try:
             with torch.no_grad():
                 outputs = model(tensor_image)
@@ -170,7 +130,6 @@ def predict(image):
                 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}")
@@ -189,7 +148,7 @@ def predict(image):
 # Gradio interface
 interface = gr.Interface(
     fn=predict,
-    inputs=gr.Image(type="pil"),  # PIL 이미지 타입으로 명시
+    inputs=gr.Image(type="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'.",