| import matplotlib.pyplot as plt | |
| import requests, validators | |
| import torch | |
| import pathlib | |
| import numpy as np | |
| from PIL import Image | |
| from transformers import DetrFeatureExtractor, DetrForSegmentation, MaskFormerImageProcessor, MaskFormerForInstanceSegmentation | |
| from transformers.models.detr.feature_extraction_detr import rgb_to_id | |
| TEST_IMAGE = Image.open(r"images/Test_Street_VisDrone.JPG") | |
| MODEL_NAME_DETR = "facebook/detr-resnet-50-panoptic" | |
| MODEL_NAME_MASKFORMER = "facebook/maskformer-swin-large-coco" | |
| DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
| ####### | |
| # Parameters | |
| ####### | |
| image = TEST_IMAGE | |
| model_name = MODEL_NAME_MASKFORMER | |
| # Starting with MaskFormer | |
| processor = MaskFormerImageProcessor.from_pretrained(model_name) | |
| model = MaskFormerForInstanceSegmentation.from_pretrained(model_name) | |
| model.to(DEVICE) | |
| # img = np.array(TEST_IMAGE) | |
| inputs = processor(images=image, return_tensors="pt") | |
| inputs.to(DEVICE) | |
| outputs = model(**inputs) |