Update tasks/image.py

tasks/image.py

@@ -35,13 +35,16 @@ model.load_state_dict(torch.load(model_path, map_location=torch.device('cpu')))
 model.eval()
 
 def preprocess(image):
+    # Ensure input image is resized to a fixed size (512, 512)
     image = image.resize((512, 512))
-    image = np.array(image)[:, :, ::-1]  # RGB to BGR
+
+    # Convert to NumPy and ensure BGR normalization
+    image = np.array(image)[:, :, ::-1]  # Convert RGB to BGR
     image = np.array(image, dtype=np.float32) / 255.0
 
-    # Convert back to PIL Image to maintain compatibility with feature extractor
-    image = Image.fromarray((image * 255).astype(np.uint8))
-    return image
+    # Return as a PIL Image for feature extractor compatibility
+    return Image.fromarray((image * 255).astype(np.uint8))
+
 
 
 def get_bounding_boxes_from_mask(mask):
@@ -145,16 +148,11 @@ async def evaluate_image(request: ImageEvaluationRequest):
         # Extract image and annotations
         image = example["image"]
 
-        original_shape = image.size
-        image = preprocess(image)
-        
+        original_shape = image.size        
         annotation = example.get("annotations", "").strip()
-        
-    
         has_smoke = len(annotation) > 0
         true_labels.append(1 if has_smoke else 0)
     
-
         if has_smoke:
             image_true_boxes = parse_boxes(annotation)
             if image_true_boxes:
@@ -165,26 +163,30 @@ async def evaluate_image(request: ImageEvaluationRequest):
             true_boxes_list.append([])
     
         # Model Inference
-        # image_input = feature_extractor(images=image, return_tensors="pt").pixel_values
-        image_input = feature_extractor(images=[image], return_tensors="pt", padding=True).pixel_values
-
+    
+        # Preprocess image
+        image = preprocess(image)
+        
+        # Ensure correct feature extraction
+        image_input = feature_extractor(images=image, return_tensors="pt").pixel_values
+        
+        # Perform inference
         with torch.no_grad():
             outputs = model(pixel_values=image_input)
             logits = outputs.logits
-    
+        
+        # Threshold and process the segmentation mask
         probabilities = torch.sigmoid(logits)
         predicted_mask = (probabilities[0, 1] > 0.30).cpu().numpy().astype(np.uint8)
-        # predicted_mask_resized = cv2.resize(predicted_mask, (512, 512), interpolation=cv2.INTER_NEAREST)
         predicted_mask_resized = cv2.resize(predicted_mask, original_shape[::-1], interpolation=cv2.INTER_NEAREST)
-
         
-        # Extract predicted bounding boxes
+        # Extract bounding boxes
         predicted_boxes = get_bounding_boxes_from_mask(predicted_mask_resized)
         pred_boxes.append(predicted_boxes)
-    
-        # Binary prediction for smoke detection
+        
+        # Smoke prediction based on bounding box presence
+        predictions.append(1 if len(predicted_boxes) > 0 else 0)    
         print(f"Prediction : {1 if len(predicted_boxes) > 0 else 0}")
-        predictions.append(1 if len(predicted_boxes) > 0 else 0)
         
 
     # Filter only valid box pairs