Update tasks/image.py

tasks/image.py

@@ -61,12 +61,14 @@ class SmokeDataset(torch.utils.data.Dataset):
         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)
-
-        return image_input, annotation
+        # Resize image and preprocess
+        image = preprocess(image)  # Apply resizing and preprocessing
 
+        # Extract features with padding set to True
+        features = feature_extractor(images=image, return_tensors="pt", padding=True)
+        
+        # Return pixel values directly as tensors
+        return features.pixel_values.squeeze(0), annotation
 
 
 def preprocess_batch(images):
@@ -184,10 +186,9 @@ async def evaluate_image(request: ImageEvaluationRequest):
     true_boxes_list = []  
 
     for batch_images, batch_annotations in dataloader:
-
-        batch_images = batch_images.to(device)  # Move to the correct device if using GPU
-
-        # Perform inference
+        
+        batch_images = batch_images.to(device)
+        
         with torch.no_grad():
             outputs = model(pixel_values=batch_images)
             logits = outputs.logits