Update tasks/image.py

tasks/image.py

@@ -6,20 +6,40 @@ from sklearn.metrics import accuracy_score
 import random
 import os
 
-from ultralytics import YOLO  # Import YOLO
-from .utils.evaluation import ImageEvaluationRequest
-from .utils.emissions import tracker, clean_emissions_data, get_space_info
+import os
+import torch
+import numpy as np
+from PIL import Image
+from transformers import MobileViTImageProcessor, MobileViTForSemanticSegmentation
+import cv2
+from tqdm import tqdm
+from dataset import WildfireSmokeDataset
+from torch.utils.data import DataLoader
 
 from dotenv import load_dotenv
 load_dotenv()
 
 router = APIRouter()
 
-DESCRIPTION = "YOLO Smoke Detection"
+DESCRIPTION = "Mobile-ViT Smoke Detection"
 ROUTE = "/image"
 
-yolo_model = YOLO("best.pt")
+model_path = "mobilevit_segmentation_full_data.pth"
+feature_extractor = MobileViTImageProcessor.from_pretrained("apple/deeplabv3-mobilevit-xx-small")
+model = MobileViTForSemanticSegmentation.from_pretrained("apple/deeplabv3-mobilevit-xx-small")
+model.load_state_dict(torch.load(model_path))
+model.eval()
 
+def get_bounding_boxes_from_mask(mask):
+    """Extract bounding boxes from a binary mask."""
+    pred_boxes = []
+    contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+    for contour in contours:
+        if len(contour) > 5:  # Ignore small/noisy contours
+            x, y, w, h = cv2.boundingRect(contour)
+            pred_boxes.append((x, y, x + w, y + h))
+    return pred_boxes
+    
 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"""
@@ -93,7 +113,7 @@ async def evaluate_image(request: ImageEvaluationRequest):
     # Split dataset
     train_test = dataset["train"].train_test_split(test_size=request.test_size, seed=request.test_seed)
     test_dataset = dataset["val"]#train_test["test"]
-    
+
     # Start tracking emissions
     tracker.start()
     tracker.start_task("inference")
@@ -126,28 +146,37 @@ async def evaluate_image(request: ImageEvaluationRequest):
         else:
             true_boxes_list.append([])
     
-        # results = yolo_model .predict(image, verbose=False)    # INFERENCE - prediction
-        results = yolo_model.predict(image)    # INFERENCE - prediction
-
-        if len(results[0].boxes):
-            pred_box = results[0].boxes.xywhn[0].cpu().numpy().tolist()
-            predictions.append(1)  
-            pred_boxes.append(pred_box)
-        else:
-            predictions.append(0)  
-            pred_boxes.append([])  
+        # Model Inference
+        image_input = feature_extractor(images=image, return_tensors="pt").pixel_values
+        with torch.no_grad():
+            outputs = model(pixel_values=image_input)
+            logits = outputs.logits
+    
+        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)
+    
+        # Extract predicted bounding boxes
+        predicted_boxes = get_bounding_boxes_from_mask(predicted_mask_resized)
+        pred_boxes.append(predicted_boxes)
+    
+        # Binary prediction for smoke detection
+        print(1 if len(predicted_boxes) > 0 else 0)
+        predictions.append(1 if len(predicted_boxes) > 0 else 0)
+        
 
+    # Filter only valid box pairs
     filtered_true_boxes_list = []
     filtered_pred_boxes = []
-
-    for true_boxes, pred_boxes_entry in zip(true_boxes_list, pred_boxes):  # Only see when annotation(s) is/are both on true label and prediction
-        if true_boxes and pred_boxes_entry:  
+    
+    for true_boxes, pred_boxes_entry in zip(true_boxes_list, pred_boxes):
+        if true_boxes and pred_boxes_entry:
             filtered_true_boxes_list.append(true_boxes)
             filtered_pred_boxes.append(pred_boxes_entry)
     
-
     true_boxes_list = filtered_true_boxes_list
     pred_boxes = filtered_pred_boxes
+
     
     #--------------------------------------------------------------------------------------------
     # YOUR MODEL INFERENCE STOPS HERE