Added title and examples

.gitignore

@@ -0,0 +1 @@
+.DS_Store

app.py

@@ -1,67 +1,65 @@
 import gradio as gr
-import torch
 import yolov5
+import os
 from transformers import pipeline
 
-pipeline = pipeline(task="image-classification", model="PranomVignesh/Police-vs-Public")
+imageClassifier = pipeline(task="image-classification",
+                           model="PranomVignesh/Police-vs-Public")
 
-# from transformers import AutoFeatureExtractor, AutoModelForImageClassification
 
-# extractor = AutoFeatureExtractor.from_pretrained("PranomVignesh/Police-vs-Public")
-# model = AutoModelForImageClassification.from_pretrained("PranomVignesh/Police-vs-Public")
-
-# Images
-# torch.hub.download_url_to_file('https://github.com/ultralytics/yolov5/raw/master/data/images/zidane.jpg', 'zidane.jpg')
-# torch.hub.download_url_to_file('https://raw.githubusercontent.com/WongKinYiu/yolov7/main/inference/images/image3.jpg', 'image3.jpg')
-
-def yolov5_inference(
-    image
-):
-    """
-    YOLOv5 inference function
-    Args:
-        image: Input image
-        model_path: Path to the model
-        image_size: Image size
-        conf_threshold: Confidence threshold
-        iou_threshold: IOU threshold
-    Returns:
-        Rendered image
-    """
+def predict(image):
     model = yolov5.load('./best.pt', device="cpu")
     results = model([image], size=224)
 
-    # outputs = model(**inputs)
-    # logits = outputs.logits
-    # probabilities = torch.softmax(logits, dim=1).tolist()[0]
+    predictions = imageClassifier(image)
+    classMappings = {
+        'police': "Police / Authorized Personnel",
+        'public': 'Unauthorized Person'
+    }
+    output = {}
+    for item in predictions:
+        output[classMappings[item['label']]] = item['score']
 
-    # classes = ['Police/Authorized Personnel', 'Public/Unauthorized Person']
+    return results.render()[0], output
 
-    # output = {name: float(prob) for name, prob in zip(classes, probabilities)}
 
-    probabilities = pipeline(image)
-    output = {p["label"]: p["score"] for p in probabilities}
+title = "Detecting Unauthorized Individuals with Firearms"
+
+examples = [
+    []
+]
 
-    return results.render()[0],output
-        
+title = "Detecting Unauthorized Individuals with Firearms"
+description = """
+  Try the examples at bottom to get started.
+"""
+examples = [[
+    os.path.join(os.path.abspath(''), './examples/sample_1.png'),
+    os.path.join(os.path.abspath(''), './examples/sample_2.png'),
+    os.path.join(os.path.abspath(''), './examples/sample_3.png'),
+    os.path.join(os.path.abspath(''), './examples/sample_4.png'),
+    os.path.join(os.path.abspath(''), './examples/sample_5.png'),
+    os.path.join(os.path.abspath(''), './examples/sample_6.png'),
+    os.path.join(os.path.abspath(''), './examples/sample_7.png'),
+    os.path.join(os.path.abspath(''), './examples/sample_8.png'),
+]]
 
-inputs = gr.Image(type="pil")
+inputs = gr.Image(type="pil", shape=(224, 224),
+                  label="Upload your image for detection")
 outputs = [
-    gr.Image(type="pil"),
-    gr.Label()
+    gr.Image(type="pil", label="Gun Detections"),
+    gr.Label(label="Class Prediction")
 ]
-title = "Detection"
-description = "YOLOv5 is a family of object detection models pretrained on COCO dataset. This model is a pip implementation of the original YOLOv5 model."
 
-# examples = [['zidane.jpg', 'yolov5s.pt', 640, 0.25, 0.45], ['image3.jpg', 'yolov5s.pt', 640, 0.25, 0.45]]
-demo_app = gr.Interface(
-    fn=yolov5_inference,
+interface = gr.Interface(
+    fn=predict,
     inputs=inputs,
     outputs=outputs,
     title=title,
-    # examples=examples,
-    # cache_examples=True,
-    # live=True,
-    # theme='huggingface',
+    examples=examples,
+    description=description,
+    cache_examples=True,
+    live=True,
+    theme='huggingface'
 )
-demo_app.launch(debug=True, enable_queue=True)
+interface.launch(debug=True, enable_queue=True)