Update app.py

app.py

@@ -1,7 +1,48 @@
 import gradio as gr
+import torch
+from torchvision import transforms as T
+import torchvision.models as models
+import requests
 
-def greet(name):
-  return "Hello " + name + "!!"
 
-iface = gr.Interface(fn=greet, inputs="text", outputs="text")
+def classify(image):
+  
+  ## preprocessing
+  # we need a transform step to normalise the pictures
+  transform = T.Compose([T.Resize(256), T.CenterCrop(224), T.ToTensor(), T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])])
+  
+  # we import the model
+  resnet34 = models.resnet34(pretrained=True)
+  
+  # normalise the image
+  image_transformed = transform(image)
+  # reshape
+  batch_image_transformed = torch.unsqueeze(image_transformed, 0)
+  # evaluation mode
+  resnet34.eval()
+  # get the predictions
+  output = resnet34(batch_image_transformed)
+  
+  ## labeling
+  # Load the file containing the 1,000 labels for the ImageNet dataset classes
+  url = hf_hub_url(repo_id="Selma/pytorch-resnet34", filename="imagenet_classes.txt")
+  response = requests.get(URL)
+  # write to a label file
+  open("labels.txt", "wb").write(response.content)
+  #extract the labels from the file
+  with open('labels.txt', "r") as f:
+      labels = [line.strip() for line in f.readlines()]
+  
+  ## predict the class
+  # Find the index (tensor) corresponding to the maximum score in the out tensor.
+  # Torch.max function can be used to find the information
+  _, index = torch.max(out, 1)
+  
+  # Find the score in terms of percentage by using torch.nn.functional.softmax function
+  # which normalizes the output to range [0,1] and multiplying by 100
+  percentage = torch.nn.functional.softmax(out, dim=1)[0] * 100
+  
+  return "The image depicts: " + labels[index[0]] + "with a score of " + percentage[index[0]].item() + "%"
+
+iface = gr.Interface(fn=classify, inputs="image", outputs="text")
 iface.launch()