Add application file

app.py

@@ -0,0 +1,154 @@
+import gradio as gr
+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
+import matplotlib.pyplot as plt
+
+
+class Detr(pl.LightningModule):
+
+     def __init__(self, lr, weight_decay):
+         super().__init__()
+         # replace COCO classification head with custom head
+         self.model = AutoModelForObjectDetection.from_pretrained("hustvl/yolos-small",
+                                                             num_labels=len(id2label),
+                                                             ignore_mismatched_sizes=True)
+         # see https://github.com/PyTorchLightning/pytorch-lightning/pull/1896
+         self.lr = lr
+         self.weight_decay = weight_decay
+
+     def forward(self, pixel_values):
+       outputs = self.model(pixel_values=pixel_values)
+
+       return outputs
+     
+     def common_step(self, batch, batch_idx):
+       pixel_values = batch["pixel_values"]
+       labels = [{k: v.to(self.device) for k, v in t.items()} for t in batch["labels"]]
+
+       outputs = self.model(pixel_values=pixel_values, labels=labels)
+
+       loss = outputs.loss
+       loss_dict = outputs.loss_dict
+
+       return loss, loss_dict
+
+     def training_step(self, batch, batch_idx):
+        loss, loss_dict = self.common_step(batch, batch_idx)     
+        # logs metrics for each training_step,
+        # and the average across the epoch
+        self.log("training_loss", loss)
+        for k,v in loss_dict.items():
+          self.log("train_" + k, v.item())
+
+        return loss
+
+     def validation_step(self, batch, batch_idx):
+        loss, loss_dict = self.common_step(batch, batch_idx)     
+        self.log("validation_loss", loss)
+        for k,v in loss_dict.items():
+          self.log("validation_" + k, v.item())
+
+        return loss
+
+     def configure_optimizers(self):
+        optimizer = torch.optim.AdamW(self.parameters(), lr=self.lr,
+                                  weight_decay=self.weight_decay)
+        
+        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 = 'fintune_traffic_object.ckpt'
+model = Detr.load_from_checkpoint(checkpoint, lr=2.5e-5, weight_decay=1e-4)
+
+model.to(device)
+model.eval()
+
+# colors for visualization
+COLORS = [[0.000, 0.447, 0.741], [0.850, 0.325, 0.098], [0.929, 0.694, 0.125],
+          [0.756, 0.794, 0.100], [0.466, 0.674, 0.188], [0.301, 0.745, 0.933],
+          [0.184, 0.494, 0.741], [0.494, 0.674, 0.556], [0.494, 0.301, 0.933],
+          [0.000, 0.325, 0.850], [0.745, 0.301, 0.188]]
+
+id2label = {1: 'person', 2: 'rider', 3: 'car', 4: 'bus', 5: 'truck', 6: 'bike', 7: 'motor', 8: 'traffic light', 9: 'traffic sign', 10: 'train'}
+
+# for output bounding box post-processing
+def box_cxcywh_to_xyxy(x):
+    x_c, y_c, w, h = x.unbind(1)
+    b = [(x_c - 0.5 * w), (y_c - 0.5 * h),
+         (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):
+    fig = plt.figure(figsize=(16,10))
+    plt.imshow(pil_img)
+    ax = plt.gca()
+    colors = COLORS * 100
+    for p, (xmin, ymin, xmax, ymax) in zip(prob, boxes.tolist()):
+        cl = p.argmax()
+        c = colors[cl]
+        ax.add_patch(plt.Rectangle((xmin, ymin), xmax - xmin, ymax - ymin,
+                                   fill=False, color=c, linewidth=2))
+        text = f'{id2label[cl.item()]}: {p[cl]:0.2f}'
+        ax.text(xmin, ymin, text, fontsize=10,
+                bbox=dict(facecolor=c, alpha=0.5))
+    plt.axis('off')
+    return Image.frombytes('RGB', fig.canvas.get_width_height(),fig.canvas.tostring_rgb())
+
+
+def generate_preds(processor, model, image):
+    inputs = processor(images=image, return_tensors="pt").to(device)
+    pixel_values = inputs.pixel_values.unsqueeze(0)
+    preds = model(pixel_values=pixel_values)
+    return preds
+
+
+def visualize_preds(image, preds, threshold=0.9):
+    # keep only predictions with confidence >= threshold
+    probas = preds.logits.softmax(-1)[0, :, :-1]
+    keep = probas.max(-1).values > threshold
+    
+    # convert predicted boxes from [0; 1] to image scales
+    bboxes_scaled = rescale_bboxes(preds.pred_boxes[0, keep].cpu(), image.size)
+
+    return plot_results(image, probas[keep], bboxes_scaled)
+
+
+def detect(img, model):
+
+    # Run inference
+    preds = generate_preds(feature_extractor, model, img)
+
+    return visualize_preds(img, preds)
+
+   
+interface = gr.Interface(
+    fn=detect,
+    inputs=[gr.Image(type="pil")], 
+    outputs=gr.Image(type="pil"),
+    # examples=[["example1.jpeg"], ["example2.jpeg"], ["example3.jpeg"]],
+    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> application on traffic object detection. ")
+
+interface.launch()