0llheaven/detr-finetuned-V2

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

# Abrir imagen desde la ruta en el equipo
url = "../../Pnephase2_Project/data/chest_xray/OTHER_ALL/00000020_001.png"
image = Image.open(url)

# Convertir la imagen a RGB si es en escala de grises
if image.mode != "RGB":
    image = image.convert("RGB")

processor = AutoImageProcessor.from_pretrained("0llheaven/detr-finetuned-V2")
model = AutoModelForObjectDetection.from_pretrained("0llheaven/detr-finetuned-V2")

inputs = processor(images=image, return_tensors="pt")
outputs = model(**inputs)

# Filtrar las predicciones con una precisión mayor a 0.9
# Define tamaño del objetivo
.target_sizes = torch.tensor([image.size[::-1]])
results = processor.post_process_object_detection(outputs, target_sizes=target_sizes)

# Dibujar cuadro alrededor del objeto detectado en la imagen
.draw = ImageDraw.Draw(image)
for result in results:
   scores = result["scores"]
   labels = result["labels"]
   boxes = result["boxes"]

for score, label, box in zip(scores, labels, boxes):
    if score > 0.9:  # Filtra solo los resultados con precisión mayor a 0.9
        box = [round(i, 2) for i in box.tolist()]
        draw.rectangle(box, outline="red", width=3)
        draw.text((box[0], box[1]), f"Pneumonia: {round(score.item(), 3)}", fill="red")  # Cambia la etiqueta a 'Pneumonia'

# Mostrar imagen
plt.figure(figsize=(10, 10))
plt.imshow(image)
plt.axis('off')
plt.show()

print(labels)