Switch to an easier to implement model

app.py

@@ -1,18 +1,14 @@
 import gradio as gr
 from PIL import Image
 import requests
-from transformers import CLIPProcessor, CLIPModel
+from transformers import pipeline, CLIPProcessor, CLIPModel
 import torch
-from ultralytics import YOLO
-from huggingface_hub import hf_hub_download
 
-# Download the YOLOv8 model file from Hugging Face Hub
-model_path = hf_hub_download(repo_id="arnabdhar/YOLOv8-Face-Detection", filename="yolov8n-face.pt")
+# 1. Use a standard transformers pipeline for object detection.
+# This model is reliable and loads directly from the Hub.
+face_detector = pipeline("object-detection", model="facebook/detr-resnet-50")
 
-# Load the face detection model from the downloaded file
-face_detector = YOLO(model_path)
-
-# Load the CLIP model and processor for image similarity
+# 2. Load the CLIP model for image similarity (this part was already working correctly).
 clip_model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32")
 clip_processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")
 
@@ -41,16 +37,16 @@ with torch.no_grad():
         animal_embeddings[name] = image_features
 
 def find_animal_lookalike(user_image):
-    # 1. Detect for a single human face using the YOLO model
-    face_detection_results = face_detector(user_image)
+    # 1. Detect people in the image. We'll use this to validate the input.
+    detections = face_detector(user_image)
     
-    # The result is a list, we check the number of detections in the first result
-    num_faces = len(face_detection_results[0].boxes)
-
-    if num_faces == 0:
-        return "No face detected. Please upload a clear photo of a single person.", None
-    if num_faces > 1:
-        return "Multiple faces detected. Please upload a photo of only one person.", None
+    # Filter for 'person' detections with a confidence score > 0.9
+    person_detections = [d for d in detections if d['label'] == 'person' and d['score'] > 0.9]
+    
+    if len(person_detections) == 0:
+        return "No person detected. Please upload a clear photo of a single person.", None
+    if len(person_detections) > 1:
+        return "Multiple people detected. Please upload a photo of only one person.", None
 
     # 2. Get the embedding of the user's face using CLIP
     with torch.no_grad():
@@ -60,7 +56,6 @@ def find_animal_lookalike(user_image):
     # 3. Calculate similarity with each animal
     similarities = {}
     for name, animal_embedding in animal_embeddings.items():
-        # Cosine similarity
         cos = torch.nn.CosineSimilarity(dim=1)
         similarity = cos(user_face_embedding, animal_embedding)
         similarities[name] = similarity.item()

requirements.txt

@@ -1,8 +1,6 @@
 gradio
 transformers
 torch
-torchvision
 Pillow
 requests
-ultralytics
-huggingface_hub
+timm