Spaces:
Runtime error
Runtime error
| 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() | |