feat: First version of app UI

app.py

@@ -0,0 +1,40 @@
+import gradio as gr
+import torch
+import os
+from torchvision import transforms
+from huggingface_hub import hf_hub_download
+
+
+REPO_ID = "gonzalocordova/DistractionDetectorCNN"
+FILENAME = "model_with_extended_dataset_resnet50_2023-03-28_10-42-07.pth"
+hf_hub_download(repo_id=REPO_ID, filename=FILENAME, output_dir=os.path.join(os.getcwd(), "model"))
+model = torch.load(os.path.join(os.getcwd(), "model"))
+model.eval()
+
+transform = transforms.Compose([
+            transforms.Resize(224),
+            transforms.ToTensor(),
+        ])
+
+
+def predict_fn(image_path, raw_output=False):
+        """
+        This function will predict the class of an image
+        :param image_path: The path of the image
+        :param raw_output: If True, it will return the raw output of the model
+        :return: Tuple (real class, predicted class, probability)
+        """
+
+        image = image.convert('RGB')
+        image = transform(image)
+        image = image.unsqueeze(0)
+        
+        with torch.no_grad():
+            output = model(image)
+        
+        probabilities = torch.exp(output)
+        top_p, top_class = probabilities.topk(1, dim=1)
+
+        return top_class.numpy()[0][0]
+
+gr.Interface(predict_fn, gr.inputs.Image(type="pil", label="Input Image"), outputs="label").launch()

requirements.txt

@@ -0,0 +1,3 @@
+gradio
+torch==2.0.0
+torchvision==0.15.1