StoneSeller commited on
Commit
7077e08
·
verified ·
1 Parent(s): b8cf886

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +7 -48
app.py CHANGED
@@ -2,37 +2,7 @@ import subprocess
2
  import sys
3
  import os
4
 
5
- # # Function to install or reinstall specific packages
6
- # def install(package):
7
- # subprocess.check_call([sys.executable, "-m", "pip", "install", "--force-reinstall", package])
8
-
9
- # # First, ensure NumPy is installed with the correct version
10
- # try:
11
- # import numpy as np
12
- # if not np.__version__.startswith("1.24"):
13
- # print("Installing compatible NumPy version...")
14
- # install("numpy==1.24.3")
15
- # except ImportError:
16
- # print("NumPy not found. Installing...")
17
- # install("numpy==1.24.3")
18
-
19
- # # Then install other dependencies
20
- # packages = {
21
- # "torch": "2.0.1",
22
- # "torchvision": "0.15.2",
23
- # "Pillow": "9.5.0",
24
- # "gradio": "3.50.2"
25
- # }
26
-
27
- # for package, version in packages.items():
28
- # try:
29
- # __import__(package.lower())
30
- # except ImportError:
31
- # print(f"Installing {package}...")
32
- # install(f"{package}=={version}")
33
-
34
-
35
- # 먼저 필요한 패키지들을 순서대로 설치
36
  def install_requirements():
37
  packages = [
38
  "numpy==1.24.3",
@@ -44,7 +14,7 @@ def install_requirements():
44
  for package in packages:
45
  subprocess.check_call([sys.executable, "-m", "pip", "install", "--force-reinstall", package])
46
 
47
- # 패키지 설치 실행
48
  install_requirements()
49
 
50
  import traceback
@@ -56,7 +26,7 @@ import torchvision.transforms as transforms
56
  from PIL import Image
57
  import gradio as gr
58
 
59
- # Define the model exactly as in training
60
  class ModifiedLargeNet(nn.Module):
61
  def __init__(self):
62
  super(ModifiedLargeNet, self).__init__()
@@ -75,7 +45,7 @@ class ModifiedLargeNet(nn.Module):
75
  x = self.fc2(x)
76
  return x
77
 
78
- # Load the trained model with error handling
79
  try:
80
  model = ModifiedLargeNet()
81
  state_dict = torch.load("modified_large_net.pt", map_location=torch.device("cpu"))
@@ -88,21 +58,17 @@ except Exception as e:
88
 
89
  transform = transforms.Compose([
90
  transforms.Resize((128, 128)),
91
- transforms.ToTensor(), # PILToTensor 대신 ToTensor 사용
92
  transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
93
  ])
94
 
95
  def custom_transform(pil_image):
96
- # PIL Image를 numpy array로 변환
97
  np_image = np.array(pil_image)
98
 
99
- # numpy array를 torch tensor로 변환 (채널 순서 변경 포함)
100
  tensor_image = torch.from_numpy(np_image.transpose((2, 0, 1))).float()
101
 
102
- # 값 범위를 [0, 1]로 정규화
103
  tensor_image = tensor_image / 255.0
104
 
105
- # ImageNet 정규화 적용
106
  normalize = transforms.Normalize(
107
  mean=[0.485, 0.456, 0.406],
108
  std=[0.229, 0.224, 0.225]
@@ -116,15 +82,12 @@ def process_image(image):
116
  return None
117
 
118
  try:
119
- # numpy array를 PIL Image로 변환
120
  if isinstance(image, np.ndarray):
121
  image = Image.fromarray(image.astype('uint8'))
122
 
123
- # RGB로 변환
124
  if image.mode != 'RGB':
125
  image = image.convert('RGB')
126
 
127
- # 크기 조정
128
  image = image.resize((128, 128), Image.Resampling.LANCZOS)
129
 
130
  print(f"Processed image size: {image.size}")
@@ -142,15 +105,13 @@ def predict(image):
142
  return {cls: 0.0 for cls in ["Rope", "Hammer", "Other"]}
143
 
144
  try:
145
- # 이미지 전처리
146
  processed_image = process_image(image)
147
  if processed_image is None:
148
  return {cls: 0.0 for cls in ["Rope", "Hammer", "Other"]}
149
 
150
  try:
151
- # 커스텀 변환 함수를 사용하여 텐서로 변환
152
  tensor_image = custom_transform(processed_image)
153
- tensor_image = tensor_image.unsqueeze(0) # 배치 차원 추가
154
 
155
  print(f"Input tensor shape: {tensor_image.shape}")
156
  print(f"Tensor dtype: {tensor_image.dtype}")
@@ -161,7 +122,6 @@ def predict(image):
161
  traceback.print_exc()
162
  return {cls: 0.0 for cls in ["Rope", "Hammer", "Other"]}
163
 
164
- # 예측 수행
165
  try:
166
  with torch.no_grad():
167
  outputs = model(tensor_image)
@@ -170,7 +130,6 @@ def predict(image):
170
  probabilities = F.softmax(outputs, dim=1)[0].cpu().numpy()
171
  print(f"Probabilities: {probabilities}")
172
 
173
- # 결과 반환
174
  classes = ["Rope", "Hammer", "Other"]
175
  results = {cls: float(prob) for cls, prob in zip(classes, probabilities)}
176
  print(f"Final results: {results}")
@@ -189,7 +148,7 @@ def predict(image):
189
  # Gradio interface
190
  interface = gr.Interface(
191
  fn=predict,
192
- inputs=gr.Image(type="pil"), # PIL 이미지 타입으로 명시
193
  outputs=gr.Label(num_top_classes=3),
194
  title="Mechanical Tools Classifier",
195
  description="Upload an image of a tool to classify it as 'Rope', 'Hammer', or 'Other'.",
 
2
  import sys
3
  import os
4
 
5
+
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
6
  def install_requirements():
7
  packages = [
8
  "numpy==1.24.3",
 
14
  for package in packages:
15
  subprocess.check_call([sys.executable, "-m", "pip", "install", "--force-reinstall", package])
16
 
17
+
18
  install_requirements()
19
 
20
  import traceback
 
26
  from PIL import Image
27
  import gradio as gr
28
 
29
+
30
  class ModifiedLargeNet(nn.Module):
31
  def __init__(self):
32
  super(ModifiedLargeNet, self).__init__()
 
45
  x = self.fc2(x)
46
  return x
47
 
48
+
49
  try:
50
  model = ModifiedLargeNet()
51
  state_dict = torch.load("modified_large_net.pt", map_location=torch.device("cpu"))
 
58
 
59
  transform = transforms.Compose([
60
  transforms.Resize((128, 128)),
61
+ transforms.ToTensor(),
62
  transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
63
  ])
64
 
65
  def custom_transform(pil_image):
 
66
  np_image = np.array(pil_image)
67
 
 
68
  tensor_image = torch.from_numpy(np_image.transpose((2, 0, 1))).float()
69
 
 
70
  tensor_image = tensor_image / 255.0
71
 
 
72
  normalize = transforms.Normalize(
73
  mean=[0.485, 0.456, 0.406],
74
  std=[0.229, 0.224, 0.225]
 
82
  return None
83
 
84
  try:
 
85
  if isinstance(image, np.ndarray):
86
  image = Image.fromarray(image.astype('uint8'))
87
 
 
88
  if image.mode != 'RGB':
89
  image = image.convert('RGB')
90
 
 
91
  image = image.resize((128, 128), Image.Resampling.LANCZOS)
92
 
93
  print(f"Processed image size: {image.size}")
 
105
  return {cls: 0.0 for cls in ["Rope", "Hammer", "Other"]}
106
 
107
  try:
 
108
  processed_image = process_image(image)
109
  if processed_image is None:
110
  return {cls: 0.0 for cls in ["Rope", "Hammer", "Other"]}
111
 
112
  try:
 
113
  tensor_image = custom_transform(processed_image)
114
+ tensor_image = tensor_image.unsqueeze(0)
115
 
116
  print(f"Input tensor shape: {tensor_image.shape}")
117
  print(f"Tensor dtype: {tensor_image.dtype}")
 
122
  traceback.print_exc()
123
  return {cls: 0.0 for cls in ["Rope", "Hammer", "Other"]}
124
 
 
125
  try:
126
  with torch.no_grad():
127
  outputs = model(tensor_image)
 
130
  probabilities = F.softmax(outputs, dim=1)[0].cpu().numpy()
131
  print(f"Probabilities: {probabilities}")
132
 
 
133
  classes = ["Rope", "Hammer", "Other"]
134
  results = {cls: float(prob) for cls, prob in zip(classes, probabilities)}
135
  print(f"Final results: {results}")
 
148
  # Gradio interface
149
  interface = gr.Interface(
150
  fn=predict,
151
+ inputs=gr.Image(type="pil"),
152
  outputs=gr.Label(num_top_classes=3),
153
  title="Mechanical Tools Classifier",
154
  description="Upload an image of a tool to classify it as 'Rope', 'Hammer', or 'Other'.",