Spaces:
Sleeping
Sleeping
Update app.py
Browse files
app.py
CHANGED
@@ -2,37 +2,7 @@ import subprocess
|
|
2 |
import sys
|
3 |
import os
|
4 |
|
5 |
-
|
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 |
-
|
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 |
-
|
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(),
|
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"),
|
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'.",
|