increased confidence scores

app.py

@@ -1,7 +1,5 @@
-# import os
 import json
 import io
-# import cv2
 import numpy as np
 import torch
 from fastapi import FastAPI, File, UploadFile, HTTPException
@@ -9,27 +7,34 @@ from PIL import Image
 from ultralytics import YOLO
 from torchvision import transforms
 
-# Import your utils (ensure geofencing_utils.py is in the same folder)
+# Import your utils
 from geofencing_utils import geofence_animal_classification
 
 app = FastAPI()
 
-# --- CONFIGURATION ---
-IMG_SIZE_DET = 1280
-IMG_SIZE_CLF = 480
-CONF_THRESHOLD = 0.2
-MD_CLASSES = {0: 'animal', 1: 'person', 2: 'vehicle'}
+# --- 1. GLOBAL CONFIGURATION ---
+class GlobalConfig:
+    # TOGGLES
+    USE_GEOFENCING = False  # Set to True to enable location-based filtering
+    
+    # THRESHOLDS
+    DET_CONF_THRESH = 0.1   # Minimum confidence to detect an object (MegaDetector)
+    CLF_CONF_THRESH = 0.5   # Minimum confidence to accept a species classification
+    
+    # CONSTANTS
+    IMG_SIZE_DET = 1280
+    IMG_SIZE_CLF = 480
+    MD_CLASSES = {0: 'animal', 1: 'person', 2: 'vehicle'}
 
 # Global Resources (Loaded on Startup)
 RESOURCES = {}
 
-
 @app.on_event("startup")
 async def load_models():
     """Loads models into memory once when the Space starts"""
     print("Loading models...")
     
-    # PATHS (Upload your models to the Space root folder)
+    # PATHS (Ensure these match your Hugging Face Space file structure)
     model_det_path = "species-space/MDV6-yolov9-e-1280.pt"
     model_clf_path = "species-space/always_crop_99710272_22x8_v12_epoch_00148.pt"
     labels_path = "species-space/always_crop_99710272_22x8_v12_epoch_00148.labels.txt"
@@ -44,7 +49,7 @@ async def load_models():
     
     # Load Transforms
     RESOURCES['transform'] = transforms.Compose([
-        transforms.Resize((IMG_SIZE_CLF, IMG_SIZE_CLF)),
+        transforms.Resize((GlobalConfig.IMG_SIZE_CLF, GlobalConfig.IMG_SIZE_CLF)),
         transforms.ToTensor()
     ])
     
@@ -53,11 +58,16 @@ async def load_models():
         lines = [line.strip() for line in f.readlines()]
         RESOURCES['labels_list'] = lines
         RESOURCES['taxonomy_map'] = {i: line for i, line in enumerate(lines)}
+    
+    # We load Geofence map even if disabled, in case you want to toggle it on later via code
+    try:
+        with open(geofence_path, 'r') as f:
+            RESOURCES['geofence_map'] = json.load(f)
+    except Exception as e:
+        print(f"Warning: Geofence file not found or invalid: {e}")
+        RESOURCES['geofence_map'] = None
         
-    with open(geofence_path, 'r') as f:
-        RESOURCES['geofence_map'] = json.load(f)
-        
-    print("Models loaded successfully.")
+    print(f"Models loaded. Geofencing is set to: {GlobalConfig.USE_GEOFENCING}")
 
 @app.post("/predict")
 async def predict(
@@ -69,14 +79,13 @@ async def predict(
     try:
         contents = await file.read()
         image_pil = Image.open(io.BytesIO(contents)).convert("RGB")
-        image_np = np.array(image_pil)
-        # Convert RGB to BGR for OpenCV/YOLO if needed, though Ultralytics handles PIL
     except Exception as e:
         raise HTTPException(status_code=400, detail="Invalid image file")
 
     # 2. Detect
     detector = RESOURCES['detector']
-    results = detector(image_pil, imgsz=IMG_SIZE_DET, conf=CONF_THRESHOLD, verbose=False)[0]
+    # USE GLOBAL CONFIG FOR DETECTOR THRESHOLD
+    results = detector(image_pil, imgsz=GlobalConfig.IMG_SIZE_DET, conf=GlobalConfig.DET_CONF_THRESH, verbose=False)[0]
     
     detections = []
     
@@ -87,7 +96,6 @@ async def predict(
     classifier = RESOURCES['classifier']
     transform = RESOURCES['transform']
     
-    # Buffers for batch classification
     crops_tensors = []
     crop_indices = [] # Index in 'detections' list
 
@@ -95,7 +103,7 @@ async def predict(
         cls_id = int(box.cls[0].item())
         conf = float(box.conf[0].item())
         xyxy = box.xyxy[0].tolist()
-        label_det = MD_CLASSES.get(cls_id, 'unknown')
+        label_det = GlobalConfig.MD_CLASSES.get(cls_id, 'unknown')
         
         det_obj = {
             "id": i + 1,
@@ -107,8 +115,6 @@ async def predict(
         if cls_id == 0: # Animal
             # Crop
             crop = image_pil.crop((int(xyxy[0]), int(xyxy[1]), int(xyxy[2]), int(xyxy[3])))
-            
-            # Transform & Permute (CHW -> HWC)
             t_crop = transform(crop).permute(1, 2, 0)
             
             crops_tensors.append(t_crop)
@@ -124,23 +130,34 @@ async def predict(
 
     # 4. Classify Animals
     if crops_tensors:
-        # Stack
-        batch = torch.stack(crops_tensors) # No .to(DEVICE) needed, default is CPU
+        batch = torch.stack(crops_tensors)
         
         with torch.no_grad():
             logits = classifier(batch)
             probs = torch.softmax(logits, dim=1)
+            # Get Top 5 for Geofencing logic
             top_conf, top_idx = torch.topk(probs, k=5, dim=1)
             
         for k in range(len(batch)):
             det_idx = crop_indices[k]
             
-            # Prepare for Geofencing
+            # Prepare data
             labels_list = RESOURCES['labels_list']
-            k_labels = [labels_list[top_idx[k][r].item()] for r in range(5) if top_idx[k][r] < len(labels_list)]
-            k_scores = [top_conf[k][r].item() for r in range(5) if top_idx[k][r] < len(labels_list)]
+            k_labels = []
+            k_scores = []
             
-            if k_labels:
+            for r in range(5):
+                idx = top_idx[k][r].item()
+                if idx < len(labels_list):
+                    k_labels.append(labels_list[idx])
+                    k_scores.append(top_conf[k][r].item())
+
+            if not k_labels:
+                continue
+
+            # --- LOGIC SWITCH BASED ON GLOBAL CONFIG ---
+            if GlobalConfig.USE_GEOFENCING and RESOURCES['geofence_map'] is not None:
+                # A. USE GEOFENCING
                 final_lbl, final_scr, src = geofence_animal_classification(
                     labels=k_labels,
                     scores=k_scores,
@@ -148,21 +165,37 @@ async def predict(
                     geofence_map=RESOURCES['geofence_map'],
                     taxonomy_map=RESOURCES['taxonomy_map']
                 )
-                
-                try: species_clean = final_lbl.split(';')[-1]
-                except: species_clean = final_lbl
-                
-                detections[det_idx].update({
-                    "species": final_lbl,
-                    "species_common": species_clean,
-                    "species_conf": float(final_scr),
-                    "source": src
-                })
+            else:
+                # B. RAW MODEL OUTPUT (Top 1)
+                final_lbl = k_labels[0]
+                final_scr = k_scores[0]
+                src = "raw_model"
+
+            # --- APPLY CLASSIFIER THRESHOLD ---
+            # If the score is too low, mark it as Unidentified/Unknown
+            if float(final_scr) < GlobalConfig.CLF_CONF_THRESH:
+                final_lbl = "Unidentified"
+                src = "low_confidence"
+
+            # Clean Label (remove scientific name if present)
+            try: species_clean = final_lbl.split(';')[-1].strip()
+            except: species_clean = final_lbl
+            
+            detections[det_idx].update({
+                "species": final_lbl,
+                "species_common": species_clean,
+                "species_conf": float(final_scr),
+                "source": src
+            })
 
     # 5. Return JSON
     return {
         "filename": file.filename,
         "sensor_id": sensor_id,
+        "config_used": {
+            "geofencing": GlobalConfig.USE_GEOFENCING,
+            "det_threshold": GlobalConfig.DET_CONF_THRESH
+        },
         "total_detections": len(detections),
         "detections": detections
     }