|
--- |
|
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']) |