gatesla commited on
Commit
3f1bb3e
·
1 Parent(s): af6274d

Working on script to test how the functions work

Browse files
Files changed (1) hide show
  1. understand.py +35 -0
understand.py ADDED
@@ -0,0 +1,35 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import matplotlib.pyplot as plt
2
+ import requests, validators
3
+ import torch
4
+ import pathlib
5
+ import numpy as np
6
+ from PIL import Image
7
+
8
+ from transformers import DetrFeatureExtractor, DetrForSegmentation, MaskFormerImageProcessor, MaskFormerForInstanceSegmentation
9
+ from transformers.models.detr.feature_extraction_detr import rgb_to_id
10
+
11
+ TEST_IMAGE = Image.open(r"images/Test_Street_VisDrone.JPG")
12
+ MODEL_NAME_DETR = "facebook/detr-resnet-50-panoptic"
13
+ MODEL_NAME_MASKFORMER = "facebook/maskformer-swin-large-coco"
14
+ DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
15
+
16
+ #######
17
+ # Parameters
18
+ #######
19
+ image = TEST_IMAGE
20
+ model_name = MODEL_NAME_MASKFORMER
21
+
22
+ # Starting with MaskFormer
23
+
24
+ processor = MaskFormerImageProcessor.from_pretrained(model_name)
25
+ model = MaskFormerForInstanceSegmentation.from_pretrained(model_name)
26
+
27
+ model.to(DEVICE)
28
+
29
+ # img = np.array(TEST_IMAGE)
30
+
31
+ inputs = processor(images=image, return_tensors="pt")
32
+ inputs.to(DEVICE)
33
+
34
+
35
+ outputs = model(**inputs)