Update app.py

app.py

@@ -1,94 +1,18 @@
-from __future__ import print_function, division
+import gradio as gr
+from transformers import pipeline
 
-import torch
-import torch.nn as nn
-import torch.optim as optim
-import torchvision
-from torchvision import datasets, models, transforms
-import os
-from torch.utils.data import DataLoader
-import numpy as np
-from tqdm import tqdm
-from timeit import default_timer as timer
-from datasets import load_dataset
-from datasets import load_metric
-import evaluate
-from train import *
-from transformers import AutoImageProcessor, MobileViTV2ForImageClassification, MobileViTV2Config
+pipeline = pipeline(task="image-classification", model="julien-c/hotdog-not-hotdog")
 
+def predict(input_img):
+    predictions = pipeline(input_img)
+    return input_img, {p["label"]: p["score"] for p in predictions} 
 
-def collate_fn(batch):
-    reviews, labels = zip(*batch)
-    labels = torch.Tensor(labels)
-    return reviews, labels
-
-
-def compute_metrics(eval_pred):
-    metric = load_metric("accuracy")
-    predictions = np.argmax(eval_pred.predictions, axis=1)
-    return metric.compute(predictions=predictions, references=eval_pred.label_ids)
-
-
-
-
-if __name__ == '__main__':
-    # Define path
-    # device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    device = torch.device("cpu")
-
-    testdir = f"data/grocery_data_small/test"
-    # Change to fit hardware
-    num_workers = 0
-    batch_size = 8
-    num_epochs = 5
-    learning_rate = 0.001
-    max_epochs_stop = 3
-    print_every = 1
-    num_classes = 8
-    # define image transformations
-    # define transforms
-
-    # Datasets from folders
-
-    data = datasets.ImageFolder(root=testdir, transform=transforms.ToTensor())
-
-
-    # Dataloader iterators, make sure to shuffle
-
-    dataloader = DataLoader(data, batch_size=batch_size, shuffle=True,
-                            num_workers=num_workers, collate_fn=collate_fn)
-
-
-    # Define the network with pretrained models
-
-    model_checkpoint = "apple/mobilevitv2-1.0-imagenet1k-256"
-
-    processor = AutoImageProcessor.from_pretrained(model_checkpoint)
-    model = MobileViTV2ForImageClassification.from_pretrained(model_checkpoint, num_labels=num_classes,
-                                                              ignore_mismatched_sizes=True)
-
-    #model = MobileViTV2ForImageClassification(MobileViTV2Config())
-    #model.classifier = nn.Linear(model.classifier.in_features, num_classes)
-    print(model.classifier)
-
-    model.load_state_dict(torch.load('best_model_mb_ft.pth'))
-    model.eval()
-    correct = 0
-    total = 0
-
-    with torch.no_grad():
-        for inputs, labels in tqdm(dataloader, desc="predicting"):
-            inputs = processor(images=inputs, return_tensors="pt")
-
-            outputs = model(**inputs)
-            logits = outputs.logits
-
-            predicted = logits.argmax(-1)
-
-            total += labels.size(0)
-            correct += (predicted == labels).sum().item()
-
-    accuracy = correct / total * 100
-
-    print(f"Predict Accuracy: {accuracy:.2f}%")
+gradio_app = gr.Interface(
+    predict,
+    inputs=gr.Image(label="Select hot dog candidate", sources=['upload', 'webcam'], type="pil"),
+    outputs=[gr.Image(label="Processed Image"), gr.Label(label="Result", num_top_classes=2)],
+    title="Hot Dog? Or Not?",
+)
 
+if __name__ == "__main__":
+    gradio_app.launch()