Update tasks/image.py

tasks/image.py

@@ -28,6 +28,25 @@ ROUTE = "/image"
 
 device = torch.device("cuda")
 
+
+def load_camera_models():
+    models = {}
+    folder = "cameras_dataset/"
+    cameras = ['brison-200', 'brison-110', 'courmettes-212', 'courmettes-160', 'brison-290', 'marguerite-29','default']
+
+    # Ensure the folder exists
+    if not os.path.exists(folder):
+        raise FileNotFoundError(f"The folder '{folder}' does not exist.")
+
+    # Iterate over files in the folder
+    for model_path in os.listdir(folder):
+        full_path = os.path.join(folder, model_path)
+        for camera in cameras:
+            if camera in model_path:
+                models[camera] = YOLO(full_path, task = 'detect')
+                break
+
+    return models
     
 def parse_boxes(annotation_string):
     """Parse multiple boxes from a single annotation string.
@@ -83,31 +102,6 @@ def compute_max_iou(true_boxes, pred_box):
     return max_iou
 
 
-class ClampTransform:
-    def __init__(self, min_val=0.0, max_val=1.0):
-        self.min_val = min_val
-        self.max_val = max_val
-
-    def __call__(self, tensor):
-        return torch.clamp(tensor, min=self.min_val, max=self.max_val)
-
-
-def collate_fn(batch):
-    images = [item['image'] for item in batch]
-    annotations = [item.get('annotations', '') for item in batch]
-
-    # Convert PIL Images to tensors
-    transform = transforms.Compose([
-        transforms.ToTensor(),
-        ClampTransform(min_val=0.0, max_val=1.0),
-        transforms.Resize((640, 640))
-    ])
-
-    images = [transform(img) for img in images]
-    images = torch.stack(images)
-    return {'image': images, 'annotations': annotations}
-
-
 def parse_boxes(annotation_string):
     """Parse multiple boxes from a single annotation string.
     Each box has 5 values: class_id, x_center, y_center, width, height"""
@@ -121,7 +115,6 @@ def parse_boxes(annotation_string):
             boxes.append(box)
     return boxes
 
-
 def compute_iou(box1, box2):
     """Compute Intersection over Union (IoU) between two YOLO format boxes."""
     # Convert YOLO format (x_center, y_center, width, height) to corners
@@ -151,7 +144,6 @@ def compute_iou(box1, box2):
     
     return intersection / (union + 1e-6)
 
-
 def compute_max_iou(true_boxes, pred_box):
     """Compute maximum IoU between a predicted box and all true boxes"""
     max_iou = 0
@@ -160,10 +152,10 @@ def compute_max_iou(true_boxes, pred_box):
         max_iou = max(max_iou, iou)
     return max_iou
 
-
-@router.post(ROUTE, tags=["Image Task"],
-              description=DESCRIPTION)
-async def evaluate_image(model_path: str = "models/yolo11s_best.pt", request: ImageEvaluationRequest = ImageEvaluationRequest()):
+# @router.post(ROUTE, tags=["Image Task"],
+#              description=DESCRIPTION)
+async def evaluate_image(model_path,request: ImageEvaluationRequest = ImageEvaluationRequest()):
+# def evaluate_image(model_path: str, request: ImageEvaluationRequest = ImageEvaluationRequest()):
     """
     Evaluate image classification and object detection for forest fire smoke.
     
@@ -182,12 +174,10 @@ async def evaluate_image(model_path: str = "models/yolo11s_best.pt", request: Im
     dataset = load_dataset(request.dataset_name, token=os.getenv("HF_TOKEN"))
     
     # Split dataset
-    train_test = dataset["train"]
+    # train_test = dataset["train"]
     test_dataset = dataset["val"]
-    if("yolo" in model_path):
-        model = YOLO(model_path, task="detect")
-    if("detr" in model_path):
-        model = RTDETR(model_path)
+
+    models = load_camera_models()
     
     # Start tracking emissions
     tracker.start()
@@ -202,15 +192,27 @@ async def evaluate_image(model_path: str = "models/yolo11s_best.pt", request: Im
     true_labels = []
     pred_boxes = []
     true_boxes_list = []  # List of lists, each inner list contains boxes for one image
+    # list of cameras
+    result_cameras = ['brison-200', 'brison-110', 'courmettes-212', 'courmettes-160', 'brison-290', 'marguerite-29']
 
-    for example in tqdm(test_dataset):
+    for example in test_dataset:
+     
         # Parse true annotation (YOLO format: class_id x_center y_center width height)
         annotation = example.get("annotations", "").strip()
         has_smoke = len(annotation) > 0
         true_labels.append(int(has_smoke))
-        
-        image=example["image"]
-        results = model(image, verbose=False)
+
+        image_path = example["image_name"]
+        image = example["image"]
+
+        # Extract camera name from the image path
+        camera = next((cam for cam in result_cameras if cam in image_path), None)
+        if camera:
+            results = models[camera](image, verbose=False,imgsz=1280)
+
+        else:
+            results = models["default"](image, verbose=False,imgsz=1280)
+            
         boxes = results[0].boxes.xywh.tolist()
 
         pred_has_smoke = len(boxes) > 0