Update tasks/image.py

tasks/image.py

@@ -7,6 +7,10 @@ import random
 import os
 
 from torch.utils.data import DataLoader
+from torch.utils.data import Dataset
+from PIL import Image
+import torch
+
 
 from ultralytics import YOLO
 from .utils.evaluation import ImageEvaluationRequest
@@ -36,22 +40,12 @@ model = MobileViTForSemanticSegmentation.from_pretrained("apple/deeplabv3-mobile
 model.load_state_dict(torch.load(model_path, map_location=torch.device('cpu')))
 model.eval()
 
-from torch.utils.data import 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):
+    def __init__(self, dataset, feature_extractor, target_size=(224, 224)):
         self.dataset = dataset
+        self.feature_extractor = feature_extractor
+        self.target_size = target_size
 
     def __len__(self):
         return len(self.dataset)
@@ -60,15 +54,34 @@ class SmokeDataset(torch.utils.data.Dataset):
         example = self.dataset[idx]
         image = example["image"]
         annotation = example.get("annotations", "").strip()
-        
-        # 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
+        # Ensure image is resized to a fixed target size using PIL
+        if isinstance(image, torch.Tensor):
+            image = Image.fromarray(image.numpy())
+        resized_image = image.resize(self.target_size, Image.ANTIALIAS)
+
+        # Process image using feature extractor
+        features = self.feature_extractor(images=resized_image, return_tensors="pt").pixel_values
+
+        return features.squeeze(0), annotation
+
+
+def collate_fn(batch):
+    images, annotations = zip(*batch)
+    images = torch.stack(images)  # Ensure batch has uniform shape
+    return images, annotations
+
+    
+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))
 
 
 def preprocess_batch(images):
@@ -178,7 +191,8 @@ async def evaluate_image(request: ImageEvaluationRequest):
     # Update the code below to replace the random baseline with your model inference
     #--------------------------------------------------------------------------------------------   
     smoke_dataset = SmokeDataset(test_dataset)
-    dataloader = DataLoader(smoke_dataset, batch_size=16, shuffle=False)
+    # dataloader = DataLoader(smoke_dataset, batch_size=16, shuffle=False)
+    dataloader = DataLoader(dataset, batch_size=8, collate_fn=collate_fn)
 
     predictions = []
     true_labels = []