Update app.py

app.py

@@ -5,47 +5,32 @@ from PIL import Image
 
 # Load the trained model
 model_path = "cifar_net.pth"
-model = torch.load(model_path, map_location=torch.device('cpu'))
-model = YourModelClass()  # Replace YourModelClass with the appropriate model class
-model.load_state_dict(state_dict)
+
+model = torch.load(model_path)
 model.eval()
 
-# Define class labels for CIFAR-10
-classes = ('plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck')
-
-def classify_image(image):
-    transform = transforms.Compose([
-        transforms.ToTensor(),
-        transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5))
-    ])
-
-    # Preprocess the input image
-    image = transform(image).unsqueeze(0)
-
-    # Perform inference with the model
-    outputs = model(image)
-    _, predicted = torch.max(outputs, 1)
-    predicted_class = classes[predicted.item()]
-
-    return predicted_class
-
-def classify_images(images):
-    return [classify_image(image) for image in images]
-
-inputs_image = gr.inputs.Image(label="Input Image", type="pil")
-outputs_image = gr.outputs.Label(label="Predicted Class")
-interface_image = gr.Interface(
-    fn=classify_images,
-    inputs=inputs_image,
-    outputs=outputs_image,
-    title="CIFAR-10 Image Classifier",
-    description="Classify images into one of the CIFAR-10 classes.",
-    examples=[
-        ['image_0.jpg'], 
-        ['image_1.jpg']
-    ],
-    allow_flagging=False
-)
-
-if __name__ == "__main__":
-    interface_image.launch()
+# Prepare the image for prediction
+image_path = 'download.jpg'
+image = Image.open(image_path)
+
+# Transform the image to match CIFAR-10 format
+transform = transforms.Compose([
+    transforms.Resize((32, 32)),
+    transforms.ToTensor(),
+    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))  # Normalize with CIFAR-10 mean and std
+])
+input_image = transform(image).unsqueeze(0)
+
+# Make predictions
+with torch.no_grad():
+    outputs = model(input_image)
+
+# Retrieve the predicted class label
+_, predicted = torch.max(outputs, 1)
+class_index = predicted.item()
+
+# Load the CIFAR-10 class labels
+classes = ['airplane', 'automobile', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck']
+
+# Print the predicted class label
+print('Predicted class label:', classes[class_index])