Add notebook file.

.gitattributes

@@ -32,3 +32,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+*.ipynb filter=lfs diff=lfs merge=lfs -text

Fine-tuning YOLOS for traffic object detection.ipynb

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c6328705e1cce9fc89a243319dc8b57997f3791298f312d2c6ac078cc8034e32
+size 15511753

README.md

@@ -1,5 +1,5 @@
 ---
-title: YOLOS Traffic Users
+title: YOLOS Traffic Object detection
 emoji: 🔥
 colorFrom: gray
 colorTo: yellow

app.py

@@ -3,21 +3,31 @@ import os
 import torch
 import pytorch_lightning as pl
 
-# torch.hub.download_url_to_file('http://images.cocodataset.org/val2017/000000039769.jpg', 'cats.jpg')
-# torch.hub.download_url_to_file('https://huggingface.co/datasets/nielsr/textcaps-sample/resolve/main/stop_sign.png', 'stop_sign.png')
-# torch.hub.download_url_to_file('https://cdn.openai.com/dall-e-2/demos/text2im/astronaut/horse/photo/0.jpg', 'astronaut.jpg')
-
-# os.system("wget https://github.com/hustvl/YOLOP/raw/main/weights/End-to-end.pth")
-
-from transformers import AutoFeatureExtractor, AutoModelForObjectDetection
-
-from PIL import Image, ImageDraw
 import cv2
 import numpy
-import matplotlib.pyplot as plt
+from transformers import AutoFeatureExtractor, AutoModelForObjectDetection
+from PIL import Image
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+feature_extractor = AutoFeatureExtractor.from_pretrained("hustvl/yolos-small", size=512, max_size=864)
 
 id2label = {1: 'person', 2: 'rider', 3: 'car', 4: 'bus', 5: 'truck', 6: 'bike', 7: 'motor', 8: 'traffic light', 9: 'traffic sign', 10: 'train'}
 
+# colors for visualization
+colors = [
+    [  0, 113, 188,],
+    [216,  82,  24,],
+    [236, 176,  31,],
+    [192, 202,  25,],
+    [118, 171,  47,],
+    [ 76, 189, 237,],
+    [ 46, 125, 188,],
+    [125, 171, 141,],
+    [125,  76, 237,],
+    [  0,  82, 216,],
+    [189,  76,  47,]]
+
 class Detr(pl.LightningModule):
 
      def __init__(self, lr, weight_decay):
@@ -71,10 +81,6 @@ class Detr(pl.LightningModule):
         return optimizer
 
 
-device = "cuda" if torch.cuda.is_available() else "cpu"
-
-feature_extractor = AutoFeatureExtractor.from_pretrained("hustvl/yolos-small", size=512, max_size=864)
-
 # Build model and load checkpoint
 checkpoint = './checkpoints/epoch=1-step=2184.ckpt'
 model_yolos = Detr.load_from_checkpoint(checkpoint, lr=2.5e-5, weight_decay=1e-4)
@@ -82,19 +88,6 @@ model_yolos = Detr.load_from_checkpoint(checkpoint, lr=2.5e-5, weight_decay=1e-4
 model_yolos.to(device)
 model_yolos.eval()
 
-# colors for visualization
-colors = [
-    [  0, 113, 188,],
-    [216,  82,  24,],
-    [236, 176,  31,],
-    [192, 202,  25,],
-    [118, 171,  47,],
-    [ 76, 189, 237,],
-    [ 46, 125, 188,],
-    [125, 171, 141,],
-    [125, 76, 237,],
-    [  0, 82, 216,],
-    [189, 76,  47,]]
 
 # for output bounding box post-processing
 def box_cxcywh_to_xyxy(x):
@@ -103,12 +96,14 @@ def box_cxcywh_to_xyxy(x):
          (x_c + 0.5 * w), (y_c + 0.5 * h)]
     return torch.stack(b, dim=1)
 
+
 def rescale_bboxes(out_bbox, size):
     img_w, img_h = size
     b = box_cxcywh_to_xyxy(out_bbox)
     b = b * torch.tensor([img_w, img_h, img_w, img_h], dtype=torch.float32)
     return b
 
+
 def plot_results(pil_img, prob, boxes):
 
     img = numpy.asarray(pil_img)
@@ -119,12 +114,8 @@ def plot_results(pil_img, prob, boxes):
         c1, c2 = (int(xmin), int(ymin)), (int(xmax), int(ymax))
 
         cv2.rectangle(img, c1, c2, c, thickness=2, lineType=cv2.LINE_AA)
-        
-        cv2.putText(img, f'{id2label[cl.item()]}: {p[cl]:0.2f}', [int(xmin), int(ymin)], cv2.FONT_HERSHEY_SIMPLEX, 0.5, c, 1)
-        # ax.text(xmin, ymin, text, fontsize=10,
-        #         bbox=dict(facecolor=c, alpha=0.5))
+        cv2.putText(img, f'{id2label[cl.item()]}: {p[cl]:0.2f}', [int(xmin), int(ymin)-5], cv2.FONT_HERSHEY_SIMPLEX, 0.7, c, 2)
     return Image.fromarray(img)
-    # return fig
 
 
 def generate_preds(processor, model, image):
@@ -145,20 +136,21 @@ def visualize_preds(image, preds, threshold=0.9):
 
 
 def detect(img):
-
     # Run inference
     preds = generate_preds(feature_extractor, model_yolos, img)
-
     return visualize_preds(img, preds)
 
+
+description = "This is a traffic object detector based on <a href='https://huggingface.co/docs/transformers/model_doc/yolos' style='text-decoration: underline' target='_blank'>YOLOS</a>. \n" + \
+    "The model can detect following targets: {1: 'person', 2: 'rider', 3: 'car', 4: 'bus', 5: 'truck', 6: 'bike', 7: 'motor', 8: 'traffic light', 9: 'traffic sign', 10: 'train'}."
+
    
 interface = gr.Interface(
     fn=detect,
     inputs=[gr.Image(type="pil")], 
     outputs=gr.Image(type="pil"),
-    # outputs = ['plot'],
     examples=[["./imgs/example1.jpg"], ["./imgs/example2.jpg"]],
     title="YOLOS for traffic object detection",
-    description="A downstream application for <a href='https://huggingface.co/docs/transformers/model_doc/yolos' style='text-decoration: underline' target='_blank'>YOLOS</a> which can performe traffic object detection. ")
+    description=description)
 
 interface.launch()