app.py

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()