Update Lib/Core.py

Lib/Core.py

@@ -1,87 +1,155 @@
+import os
+import sys
+sys.path.append(os.getcwd())
+
 import cv2
 import numpy as np
 import torch
 from matplotlib import pyplot as plt
 from ultralytics import YOLO
-from ultralytics.engine.results import Masks
+from Lib.Consts import LABELS, COLOR_MAP, COLOR_MAP_RGB
+
+
+DEVICE = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+
 
 class CablePoleSegmentation():
-    def __init__(self, MODEL_PATH=None, retina_mask=False):
-        if not MODEL_PATH:
-            MODEL_PATH = "./weight/yolov9c-cable-seg.pt"
+    def __init__(self, model_path=None, retina_mask=False):
+        if not model_path:
+            model_path = "./weight/yolov9c-cable-seg.pt"
         self._RetinaMask=retina_mask
-        self.Model = YOLO(MODEL_PATH)  # load a custom model
+        self.Model = None
+        self.PrepareModel(model_path)
 
 
-    def RescaleTheMask(self, orijinal_image, masks):
-        _masks = []
-        for contour in masks:
-            b_mask = np.zeros(orijinal_image.shape[:2], np.uint8)
-            contour = contour.astype(np.int32)
-            contour = contour.reshape(-1, 1, 2)
-            mask = cv2.drawContours(b_mask, [contour], -1, (1, 1, 1), cv2.FILLED)
-            _masks += [mask]
-        return _masks
+    def PrepareModel(self, model_path):
+        self.Model = YOLO(model_path)
+        self.Model.fuse()
 
 
-    def Process(self, image):
-        with torch.no_grad():
-            results = self.Model(
-                image,
-                save=False,
-                show_boxes=False,
-                project="./result/",
-                conf=0.5,
-                retina_masks=self._RetinaMask,
-                stream=True
-            )
+    def ScaleMasks(self, masks: torch.Tensor, shape: tuple) -> torch.Tensor:
+        masks = masks.unsqueeze(0)
+        interpolatedMask:torch.Tensor = torch.nn.functional.interpolate(masks, shape, mode="nearest")
+        interpolatedMask = interpolatedMask.squeeze(0)
+        return interpolatedMask
+    
+
+    def ParseResults(self, results, threshold=0.5, scale_masks=True):
+        batches = []
+        
+        SCORES = torch.Tensor([]).to(DEVICE)
+        CLASSES = torch.Tensor([]).to(DEVICE)
+        MASKS = torch.Tensor([]).to(DEVICE)
+        BOXES = torch.Tensor([]).to(DEVICE)
 
         with torch.no_grad():
             for result in results:
-                maskCountours = result.masks.xy
-                boxes = result.boxes.xyxy.int().cpu().numpy()
-                classes = result.boxes.cls.cpu().numpy()
+                original_shape = result.orig_shape
+                _scores = result.boxes.conf        # 7
+                _classes = result.boxes.cls         # 7
+                _masks = result.masks.data         # 7, 480, 640
+                _boxes = result.boxes.xyxy         # 7, 4
+                
+                # Threshold Filter
+                conditions = _scores > threshold
+                SCORES = torch.cat((SCORES, _scores[conditions]), dim=0)
+                CLASSES = torch.cat((CLASSES, _classes[conditions]), dim=0)
+                BOXES = torch.cat((BOXES, _boxes[conditions]), dim=0)
+                mask = _masks[conditions]
+
+                if mask.shape[0] == 0:
+                    continue
+
+                if scale_masks:
+                    mask = self.ScaleMasks(mask, original_shape[:2])
+
+                MASKS = torch.cat((MASKS, mask), dim=0)
+                
+                batches += [(SCORES, CLASSES, MASKS, BOXES)]
 
-            rescaledMasks =  self.RescaleTheMask(image, maskCountours)
-            return rescaledMasks, boxes, classes, result.plot()
+        return batches
 
 
+    def DrawResults(self, image, scores: torch.Tensor, classes: torch.Tensor, masks: torch.Tensor, boxes: torch.Tensor, labels:dict=LABELS, class_filter:list=None):
+        _image = np.array(image).copy()
+        _image = cv2.cvtColor(_image, cv2.COLOR_BGR2RGB)
+        maskCanvas = np.zeros_like(_image)
 
-    def PlotResults(self, masks, boxes, classes, original_image, result_image, mask, cable_mask):
-        fig, axs = plt.subplots(nrows=2, ncols=2, figsize=(27,15))
-        axs[0][0].imshow(original_image)
-        axs[0][0].set_title("Orijinal Görüntü")
 
-        axs[0][1].imshow(mask)
-        axs[0][1].set_title("Segmentasyon Maskesi")
+        with torch.no_grad():
+            scores = scores.cpu().numpy()
+            classes = classes.cpu().numpy().astype(np.int32)
+            masks = masks.cpu().numpy()
+            boxes = boxes.cpu().numpy()
+            colors = list(COLOR_MAP_RGB.values())
+
+        for score, cls, mask, box in zip(scores, classes, masks, boxes):
+            label = labels[cls]
+            _color = colors[cls]
+
+            if class_filter and cls not in class_filter:
+                continue
+
+            box = box.astype(np.int32)
+            mask = (cv2.cvtColor(mask, cv2.COLOR_GRAY2BGR)*_color).astype(np.uint8)
+            maskCanvas = cv2.addWeighted(maskCanvas, 1.0, mask, 1.0, 0)
+            maskCanvas = cv2.rectangle(maskCanvas, (box[0], box[1]), (box[2], box[3]), color=_color, thickness=5)  # Red color for bounding box
+            maskCanvas = cv2.putText(maskCanvas, f"{label} : {score:.2f}", (box[0], box[1] - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, color=_color, thickness=2)
+        
+        canvas = cv2.addWeighted(_image, 1.0, maskCanvas.astype(np.uint8), 0.5, 0)
+        return canvas, maskCanvas
 
+
+    def Process(self, image, model_threshold=0.6, overall_threshold=0.6, iou=0.7,  class_filter:list=None):
         
-        cv2.imwrite("cable_mask.png", cable_mask)
-        axs[1][0].imshow(cable_mask)
-        axs[1][0].set_title("Seçilen")
+        with torch.no_grad():
+            results = self.Model(
+                image,
+                save=False,
+                show_boxes=False,
+                project="./inference/",
+                conf=model_threshold,
+                iou=iou,
+                retina_masks=False,
+                stream=True,
+                classes=class_filter,
+                device=DEVICE
+            )
+
+            batches = self.ParseResults(results, threshold=overall_threshold, scale_masks=True)
+
+            return batches
+            
 
-        axs[1][1].imshow(result_image)
-        axs[1][1].set_title("Sonuç")
-        plt.show()
 
 
 
 if "__main__" == __name__:
-    test = "data/16_3450.png"
+    test = "data/DJI_20240905091530_0003_W.JPG"
     image = cv2.imread(test)
-    model = CablePoleSegmentation(retina_mask=True)
-    masks, boxes, classes, result_plot = model.Process(image)
+    model = CablePoleSegmentation(retina_mask=False)
+    batches = model.Process(image)
 
+    if len(batches) == 0:
+        exit()
 
-    fig, axs = plt.subplots(nrows=2, ncols=2, figsize=(27,15))
-    axs[0][0].imshow(image)
-    axs[0][0].set_title("Orijinal Görüntü")
-    axs[1][1].imshow(np.any(masks, axis=0))
-    axs[1][1].set_title("Sonuç")
-    plt.show()
+    scores, classes, masks, boxes = batches[0] # First
+    canvas, mask = model.DrawResults(image, scores, classes, masks, boxes, class_filter=None)
+    print(canvas.shape)
 
 
-    # model.PlotResults(*model.Process(image))
+    #! Plot
+    fig, axs = plt.subplots(1, 3, figsize=(27, 15))
+    axs[0].imshow(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
+    axs[0].set_title("Orijinal Görüntü")
 
+    axs[1].imshow(mask)
+    axs[1].set_title("Segmentasyon Maskesi")
+
+    axs[2].imshow(canvas)
+    axs[2].set_title("Sonuç")
+
+    plt.tight_layout()
+    plt.show()