Assignment1 / app.py
StoneSeller's picture
Update app.py
d16040b verified
raw
history blame
4.69 kB
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}")
import traceback
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
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__()
self.name = "modified_large"
self.conv1 = nn.Conv2d(3, 5, 5)
self.pool = nn.MaxPool2d(2, 2)
self.conv2 = nn.Conv2d(5, 10, 5)
self.fc1 = nn.Linear(10 * 29 * 29, 32)
self.fc2 = nn.Linear(32, 3)
def forward(self, x):
x = self.pool(F.relu(self.conv1(x)))
x = self.pool(F.relu(self.conv2(x)))
x = x.view(-1, 10 * 29 * 29)
x = F.relu(self.fc1(x))
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"))
model.load_state_dict(state_dict)
print("Model loaded successfully")
model.eval()
except Exception as e:
print(f"Error loading model: {str(e)}")
traceback.print_exc()
# Define image transformation pipeline
transform = transforms.Compose([
transforms.Resize((128, 128)),
transforms.PILToTensor(), # Changed from ToTensor()
transforms.ConvertImageDtype(torch.float32),
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])
def process_image(image):
if image is None:
return None
try:
# Convert numpy array to PIL Image
if isinstance(image, np.ndarray):
image = Image.fromarray(image.astype('uint8'))
# Convert to RGB if necessary
if image.mode != 'RGB':
image = image.convert('RGB')
print(f"Processed image size: {image.size}")
print(f"Processed image mode: {image.mode}")
return image
except Exception as e:
print(f"Error in process_image: {str(e)}")
traceback.print_exc()
return None
def predict(image):
if image is None:
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 using torchvision transforms
try:
tensor_image = transform(processed_image).unsqueeze(0)
print(f"Input tensor shape: {tensor_image.shape}")
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}")
probabilities = F.softmax(outputs, dim=1)[0].cpu().numpy()
print(f"Probabilities: {probabilities}")
# 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()
return {cls: 0.0 for cls in ["Rope", "Hammer", "Other"]}
# Gradio interface
interface = gr.Interface(
fn=predict,
inputs=gr.Image(),
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'.",
)
# Launch the interface
if __name__ == "__main__":
interface.launch()