Maicol2001's picture
Update app.py
d386553 verified
from transformers import DetrForObjectDetection, DetrFeatureExtractor
from PIL import Image
import requests
# Cargar el modelo y el extractor de características pre-entrenados
model = DetrForObjectDetection.from_pretrained('facebook/detr-resnet-50')
feature_extractor = DetrFeatureExtractor.from_pretrained('facebook/detr-resnet-50')
def object_exists(image_url):
# Cargar la imagen desde la URL
image = Image.open(requests.get(image_url, stream=True).raw)
# Preparar las entradas
inputs = feature_extractor(images=image, return_tensors="pt")
# Pasar por el modelo
outputs = model(**inputs)
# Obtener las probabilidades de clase
logits = outputs.logits
probas = logits.softmax(-1)[0, :, :-1]
# Establecer un umbral de confianza
threshold = 0.9
# Encontrar índices de clases con probabilidad por encima del umbral
keep = probas.max(-1).values > threshold
if keep.sum() > 0:
print("Existe un objeto en la imagen.")
else:
print("No existe un objeto en la imagen.")
if __name__ == "__main__":
# Reemplazar con la URL de tu imagen subida a Hugging Face
image_url = 'https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/beach.png'
object_exists(image_url)