initial commit

.gitignore

@@ -0,0 +1,3 @@
+.DS_Store
+/venv/
+__pycache__

app.py

@@ -0,0 +1,22 @@
+import gradio as gr
+from image_classifier import ImageClassifier
+from PIL import Image
+
+classifier = ImageClassifier('models/trained_model.pth')
+examples = [
+    'examples/0ab9373f-4d97-456a-9b9d-60b4d05d102e.jpg',
+    'examples/093e6044-5809-44d6-a93f-f06a6702f20d.jpg',
+    'examples/332b50b9-ab27-42ac-9118-ae32b3458d97.jpg',
+    'examples/357fbd95-ff50-4d90-9487-531595f02a9e.jpg',
+    'examples/b100a944-0e67-4e4d-9ccc-c6c63fb99c8e.jpg'
+]
+
+def predict(image):
+    image = Image.fromarray(image.astype('uint8'), 'RGB')
+    return classifier.classify_image(image)
+
+
+image = gr.inputs.Image(shape=(192,192))
+label = gr.outputs.Label()
+intf = gr.Interface(fn=predict, inputs=image, outputs=label, examples=examples)
+intf.launch(inline=False)

image_classifier.py

@@ -0,0 +1,34 @@
+import torch
+from torchvision import models, transforms
+from PIL import Image
+
+CLASS_NAMES = ['apple', 'bread', 'fried_chicken', 'hamburger', 'pizza', 'popcorn', 'salad', 'steak', 'taco']
+
+class ImageClassifier:
+
+    def __init__(self, model_path, device='cpu'):
+        self.transform = transforms.Compose([
+            transforms.Resize(256),
+            transforms.CenterCrop(224),
+            transforms.ToTensor(),
+            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
+        ])
+
+        self.model = models.resnet50(weights='ResNet50_Weights.DEFAULT')
+
+        # Adjust the last layer to match the number of classes
+        num_ftrs = self.model.fc.in_features
+        self.model.fc = torch.nn.Linear(num_ftrs, len(CLASS_NAMES))
+
+        # Load the saved model
+        self.model.load_state_dict(torch.load(model_path, map_location=torch.device(device)))
+        self.model.eval()  # Set the model to evaluation mode
+
+    def classify_image(self, image):
+        image = self.transform(image).unsqueeze(0)  # Add batch dimension
+
+        # Perform inference
+        with torch.no_grad():
+            output = self.model(image)
+            _, predicted = torch.max(output, 1)
+            return CLASS_NAMES[predicted.item()]

models/trained_model.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6e39765d1161e22f416790ca8a8857141eaed8752565fe009fd91c4583b9f7bb
+size 94405309