File size: 1,508 Bytes
cdfca96
 
 
 
ba1484f
cdfca96
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
---
license: apache-2.0
---
-------------------------HOW TO USE---------------------

from transformers import AutoProcessor, AutoModelForObjectDetection
from PIL import Image
import torch
import matplotlib.pyplot as plt

COLORS = [[0.000, 0.447, 0.741], [0.850, 0.325, 0.098], [0.929, 0.694, 0.125],
          [0.494, 0.184, 0.556], [0.466, 0.674, 0.188], [0.301, 0.745, 0.933]]
def plot_results(pil_img, scores, labels, boxes):
    plt.figure(figsize=(16,10))
    plt.imshow(pil_img)
    ax = plt.gca()
    colors = COLORS * 100
    for score, label, (xmin, ymin, xmax, ymax),c  in zip(scores.tolist(), labels.tolist(), boxes.tolist(), colors):
        ax.add_patch(plt.Rectangle((xmin, ymin), xmax - xmin, ymax - ymin,
                                   fill=False, color=c, linewidth=3))
        text = f'{model.config.id2label[label]}: {score:0.2f}'
        ax.text(xmin, ymin, text, fontsize=15,
                bbox=dict(facecolor='yellow', alpha=0.5))
    plt.axis('off')
    plt.show()

processor = AutoProcessor.from_pretrained('Lam-Hung/orange_detetion_model')
model = AutoModelForObjectDetection.from_pretrained('Lam-Hung/orange_detetion_model').to('cuda')

image = Image.open('/content/C59_59.jpeg')
inputs = processor(image, return_tensors="pt").to('cuda')
with torch.no_grad():
    outputs = model(**inputs)
results = processor.post_process_object_detection(outputs, threshold=0.3, target_sizes=[image.size[::-1]])[0]
plot_results(image, results['scores'], results['labels'], results['boxes'])