init

app.py

@@ -0,0 +1,46 @@
+import torch
+import torch.nn as nn
+from torchvision import models, transforms
+from PIL import Image
+import gradio as gr
+
+# Load your resnet18 model from Hugging Face
+model = models.resnet18()
+model.fc = nn.Linear(model.fc.in_features, 4)  # Assuming 4 classes
+checkpoint = torch.hub.load_state_dict_from_url(
+    'https://huggingface.co/wandikafp/resnet18-tom-and-jerry-classifier/resolve/main/pytorch_model.bin',
+    map_location=torch.device('cpu')
+)
+model.load_state_dict(checkpoint)
+model.eval()
+
+# Define image transformations
+transform = transforms.Compose([
+    transforms.Resize((224, 224)),
+    transforms.ToTensor(),
+    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+])
+
+# Define a prediction function
+def classify_image(image):
+    image = Image.fromarray(image)  # Convert to PIL image
+    image = transform(image).unsqueeze(0)  # Preprocess the image
+
+    with torch.no_grad():
+        outputs = model(image)
+        _, predicted = torch.max(outputs, 1)
+
+    labels = ['tom', 'jerry', 'tom_jerry_0', 'tom_jerry_1']
+    return labels[predicted.item()]
+
+# Create Gradio interface
+interface = gr.Interface(
+    fn=classify_image, 
+    inputs="image", 
+    outputs="label",
+    title="Tom and Jerry Classifier",
+    description="Classify images as 'tom', 'jerry', 'tom_jerry_0', or 'tom_jerry_1'."
+)
+
+# Launch the Gradio app
+interface.launch()

requirements.txt

@@ -0,0 +1,3 @@
+torch
+torchvision
+gradio