Update app.py

app.py

@@ -1,8 +1,13 @@
 import os
-from huggingface_hub import Repository
-import numpy as np
+import torch
+from torch.utils.data import Dataset, DataLoader
 import pandas as pd
-import matplotlib.pyplot as plt
+import numpy as np
+import joblib
+from PIL import Image
+from torchvision import transforms
+from sklearn.preprocessing import LabelEncoder
+from gradio import Interface, Image, Label
 
 # Retrieve the token from the environment variables
 token = os.environ.get("token")
@@ -15,21 +20,6 @@ repo = Repository(
 )
 repo.git_pull()
 
-import torch
-from torch.utils.data import Dataset, DataLoader
-import pandas as pd
-import numpy as np
-import io
-import joblib
-import requests
-from tqdm import tqdm
-from PIL import Image
-from torchvision import transforms
-from sklearn.preprocessing import LabelEncoder
-from sklearn.model_selection import train_test_split
-from torchvision import models
-import gradio as gr
-
 device = 'cpu'
 le = LabelEncoder()
 le = joblib.load("SVD/le.gz")
@@ -45,7 +35,6 @@ class ModelPre(torch.nn.Module):
             torch.nn.ReLU(),
             torch.nn.Linear(in_features=512,out_features=len_classes),
         )
-        # Freeze all layers
 
     def forward(self, data):
         return self.embedding(data)
@@ -55,9 +44,6 @@ model = torch.load("SVD/GeoG.pth", map_location=torch.device(device))
 modelm = ModelPre()
 modelm.load_state_dict(model['model'])
 
-import warnings
-warnings.filterwarnings("ignore", category=RuntimeWarning, module="multiprocessing.popen_fork")
-
 cmp = transforms.Compose([
     transforms.ToTensor(),
     transforms.Resize(size=(224, 224), antialias=True),
@@ -76,29 +62,17 @@ def predict(input_img):
         results = {top_10_predictions[i]: float(top_10_probabilities[i]) for i in range(10)}
         return results
 
-def create_bar_plot(predictions):
-    data = pd.DataFrame(list(predictions.items()), columns=["Location", "Probability"])
-    max_prob = data["Probability"].max()
-    return gr.BarPlot(
-        data,
-        x="Location",
-        y="Probability",
-        title="Top 10 Predictions with Probabilities",
-        tooltip=["Location", "Probability"],
-        y_lim=[0, max_prob],
-        width=800,  # Set the width of the plot
-        height=600  # Set the height of the plot
-    )
+def create_label_output(predictions):
+    return predictions
 
 def predict_and_plot(input_img):
     predictions = predict(input_img)
-    return create_bar_plot(predictions)
-    
-    
-gradio_app = gr.Interface(
+    return create_label_output(predictions)
+
+gradio_app = Interface(
     fn=predict_and_plot,
-    inputs=gr.Image(label="Upload an Image", type="pil"),
-    outputs=gr.BarPlot(),
+    inputs=Image(label="Upload an Image", type="pil"),
+    outputs=Label(num_top_classes=10),
     title="Predict the Location of this Image"
 )