Update app.py

app.py

@@ -10,7 +10,7 @@ import torchvision.transforms as T
 from PIL import Image
 import gradio as gr
 from datetime import datetime
-
+from ultralytics import YOLO
 import models
 
 print(f"Is CUDA available: {torch.cuda.is_available()}")
@@ -66,10 +66,45 @@ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 linear_model_name = 'linear_model.pt'
 classify_model = models.LinearClassifier(input_dim=768, output_dim=num_classes)
 classify_model.load_state_dict(torch.load(linear_model_name))
-
+detect_model = YOLO('yolov8m_2023-10-23_best.pt')
 k = 5
 
+def detect(image):
+    results = detect_model(image, conf=0.1)
+    # Get the current time
+    current_time = datetime.now()
+    # Format the current time as a string
+    formatted_time = current_time.strftime("%Y-%m-%d %H:%M:%S")
+    print(formatted_time)
+    try:
+        results = results[0].boxes.xyxy[0].cpu().numpy()
+        top = int(results[1])
+        left = int(results[0])
+        width = int(results[2] - results[0])
+        height = int(results[3] - results[1])
+        return {
+            "top": top,
+            "left": left,
+            "width": width,
+            "height": height
+        }
+    except:
+        return {
+            "top": 0,
+            "left": 0,
+            "width": 0,
+            "height": 0
+        }
+
 def classify(image):
+    detection = detect(image)
+    
+    if detection["top"] == 0 and detection["left"] == 0 and detection["width"] == 0 and detection["height"] == 0:
+        return {}
+    # Crop the image
+    image = image.crop((detection['left'], detection['top'], detection['left'] + detection['width'], detection['top'] + detection['height']))
+    
+    # Perform the embedding search
     embedding = extract_embedding(image)
     embedding = embedding['embedding']
     output = classify_model(torch.Tensor(embedding).to(device))