KabeerAmjad
commited on
Commit
•
75a5b88
1
Parent(s):
f63495a
Update app.py
Browse files
app.py
CHANGED
@@ -1,66 +1,62 @@
|
|
1 |
-
import gradio as gr
|
2 |
import torch
|
3 |
-
|
4 |
-
|
5 |
from PIL import Image
|
6 |
-
import
|
7 |
-
|
8 |
-
#
|
9 |
-
|
10 |
-
|
11 |
-
|
12 |
-
#
|
13 |
-
|
14 |
-
|
15 |
-
|
16 |
-
|
17 |
-
|
18 |
-
|
19 |
-
|
20 |
-
|
21 |
-
|
22 |
-
|
23 |
-
|
24 |
-
|
25 |
-
model.eval() # Set the model to evaluation mode
|
26 |
-
|
27 |
-
# Define the same preprocessing used during training
|
28 |
-
transform = transforms.Compose([
|
29 |
-
transforms.Resize((224, 224)),
|
30 |
transforms.ToTensor(),
|
31 |
-
transforms.Normalize(
|
|
|
|
|
|
|
32 |
])
|
33 |
|
34 |
-
#
|
35 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
36 |
# Preprocess the image
|
37 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
38 |
|
39 |
-
#
|
40 |
with torch.no_grad():
|
41 |
-
|
42 |
-
probs = torch.softmax(outputs, dim=-1)
|
43 |
|
44 |
-
# Get the
|
45 |
-
|
46 |
-
|
47 |
-
# Map label index to the actual class name
|
48 |
-
label_mapping = {
|
49 |
-
0: "apple_pie", 1: "cheesecake", 2: "chicken_curry", 3: "french_fries",
|
50 |
-
4: "fried_rice", 5: "hamburger", 6: "hot_dog", 7: "ice_cream",
|
51 |
-
8: "omelette", 9: "pizza", 10: "sushi"
|
52 |
-
}
|
53 |
-
return label_mapping[top_label]
|
54 |
-
|
55 |
-
# Create the Gradio interface
|
56 |
-
iface = gr.Interface(
|
57 |
-
fn=classify_image,
|
58 |
-
inputs=gr.Image(type="pil"),
|
59 |
-
outputs="text",
|
60 |
-
title="Food Image Classification",
|
61 |
-
description="Upload an image to classify if it’s an apple pie, etc."
|
62 |
-
)
|
63 |
|
64 |
-
|
65 |
-
iface.launch()
|
66 |
|
|
|
|
|
|
|
|
|
|
|
|
1 |
import torch
|
2 |
+
import torchvision.transforms as transforms
|
3 |
+
import torchvision.models as models
|
4 |
from PIL import Image
|
5 |
+
import json
|
6 |
+
|
7 |
+
# Load the model with updated weights parameter
|
8 |
+
model = models.resnet50(weights=models.ResNet50_Weights.DEFAULT)
|
9 |
+
model.eval() # Set model to evaluation mode
|
10 |
+
|
11 |
+
# Load the model's custom state_dict
|
12 |
+
model_path = 'path_to_your_model_file.pth'
|
13 |
+
try:
|
14 |
+
state_dict = torch.load(model_path, map_location=torch.device('cpu'))
|
15 |
+
model.load_state_dict(state_dict)
|
16 |
+
except RuntimeError as e:
|
17 |
+
print("Error loading state_dict:", e)
|
18 |
+
print("Ensure that the saved model architecture matches ResNet50.")
|
19 |
+
|
20 |
+
# Define the image transformations
|
21 |
+
preprocess = transforms.Compose([
|
22 |
+
transforms.Resize(256),
|
23 |
+
transforms.CenterCrop(224),
|
|
|
|
|
|
|
|
|
|
|
24 |
transforms.ToTensor(),
|
25 |
+
transforms.Normalize(
|
26 |
+
mean=[0.485, 0.456, 0.406],
|
27 |
+
std=[0.229, 0.224, 0.225],
|
28 |
+
),
|
29 |
])
|
30 |
|
31 |
+
# Load labels
|
32 |
+
with open("imagenet_classes.json") as f:
|
33 |
+
labels = json.load(f)
|
34 |
+
|
35 |
+
# Function to predict image class
|
36 |
+
def predict(image_path):
|
37 |
+
# Open the image file
|
38 |
+
input_image = Image.open(image_path).convert("RGB")
|
39 |
+
|
40 |
# Preprocess the image
|
41 |
+
input_tensor = preprocess(input_image)
|
42 |
+
input_batch = input_tensor.unsqueeze(0) # Add batch dimension
|
43 |
+
|
44 |
+
# Check if a GPU is available and move the input and model to GPU
|
45 |
+
if torch.cuda.is_available():
|
46 |
+
input_batch = input_batch.to('cuda')
|
47 |
+
model.to('cuda')
|
48 |
|
49 |
+
# Perform inference
|
50 |
with torch.no_grad():
|
51 |
+
output = model(input_batch)
|
|
|
52 |
|
53 |
+
# Get the predicted class with the highest score
|
54 |
+
_, predicted_idx = torch.max(output, 1)
|
55 |
+
predicted_class = labels[str(predicted_idx.item())]
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
56 |
|
57 |
+
return predicted_class
|
|
|
58 |
|
59 |
+
# Example usage
|
60 |
+
image_path = 'path_to_your_image.jpg'
|
61 |
+
predicted_class = predict(image_path)
|
62 |
+
print(f"Predicted class: {predicted_class}")
|