Update tasks/image.py

tasks/image.py

@@ -38,7 +38,18 @@ model.eval()
 
 from torch.utils.data import Dataset
 
-class SmokeDataset(Dataset):
+def preprocess(image):
+    # Ensure input image is resized to a fixed size (512, 512)
+    image = image.resize((512, 512))
+
+    # 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
+
+    # Return as a PIL Image for feature extractor compatibility
+    return Image.fromarray((image * 255).astype(np.uint8))
+
+class SmokeDataset(torch.utils.data.Dataset):
     def __init__(self, dataset):
         self.dataset = dataset
 
@@ -49,26 +60,14 @@ class SmokeDataset(Dataset):
         example = self.dataset[idx]
         image = example["image"]
         annotation = example.get("annotations", "").strip()
+        
+        # Preprocess and extract features directly within the dataset
+        image = preprocess(image)  # Apply resizing and other preprocessing
+        image_input = feature_extractor(images=image, return_tensors="pt").pixel_values.squeeze(0)
 
-        # Resize and preprocess the image
-        image = image.resize((512, 512))
-        image = np.array(image)[:, :, ::-1]  # Convert RGB to BGR
-        image = np.array(image, dtype=np.float32) / 255.0
-
-        # Return both the preprocessed image tensor and annotation
-        return torch.tensor(image, dtype=torch.float32).permute(2, 0, 1), annotation
-
-
-def preprocess(image):
-    # Ensure input image is resized to a fixed size (512, 512)
-    image = image.resize((512, 512))
+        return image_input, annotation
 
-    # 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
 
-    # Return as a PIL Image for feature extractor compatibility
-    return Image.fromarray((image * 255).astype(np.uint8))
 
 def preprocess_batch(images):
     """
@@ -185,12 +184,12 @@ async def evaluate_image(request: ImageEvaluationRequest):
     true_boxes_list = []  
 
     for batch_images, batch_annotations in dataloader:
-        # image_inputs = feature_extractor(images=batch_images, return_tensors="pt", padding=True).pixel_values
-        image_inputs = feature_extractor(images=[img.permute(1, 2, 0).numpy() for img in batch_images], return_tensors="pt", padding=True).pixel_values
+
+        batch_images = batch_images.to(device)  # Move to the correct device if using GPU
 
         # Perform inference
         with torch.no_grad():
-            outputs = model(pixel_values=image_inputs)
+            outputs = model(pixel_values=batch_images)
             logits = outputs.logits
     
         probabilities = torch.sigmoid(logits)