Update Predict.py

Predict.py

@@ -1,57 +1,153 @@
 import cv2
-from matplotlib import pyplot as plt
 import numpy as np
-from ultralytics import YOLO
 import torch
+from matplotlib import pyplot as plt
+from ultralytics import YOLO
+
+
+DEVICE = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 
 # Data
-test01 = "data/16_3450.png"
-test_image = test01
+test_image = "data/DJI_20240905091530_0003_W.JPG"
+
+LABELS = {0: "Boş", 1: "Çelik Direkler", 2: "Kafes Kule", 3: "Kablo", 4: "Ahşap Kule"}
+colorMap = {"Boş":"#ffffff", "Çelik Direkler":"#0000ff", "Kafes Kule":"#ff0000", "Kablo":"#00ff00", "Ahşap Kule":"#ff0000"}
 
 # Load a model
-model = YOLO("weight/yolov9c-cable-seg.pt")  # load a custom model
+model = YOLO("Weight/yolov9c-cable-seg.pt")  # load a custom model
+model.fuse()
+
+
+def ParseResults(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:
+            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 scale_masks:
+                mask = ScaleMasks(mask, original_shape[:2])
+
+            MASKS = torch.cat((MASKS, mask), dim=0)
+            
+            batches += [(SCORES, CLASSES, MASKS, BOXES)]
+
+    return batches
+ 
+
+def ScaleMasks(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 DrawResults(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)
+
+
+    with torch.no_grad():
+        scores = scores.cpu().numpy()
+        classes = classes.cpu().numpy().astype(np.int32)
+        masks = masks.cpu().numpy()
+        boxes = boxes.cpu().numpy()
+    
+    for score, cls, mask, box in zip(scores, classes, masks, boxes):
+        label = labels[cls]
+
+        if class_filter and cls not in class_filter:
+            continue
+
+        box = box.astype(np.int32)
+        mask = cv2.cvtColor(mask*255, cv2.COLOR_GRAY2BGR).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=(255, 0, 0), thickness=3)  # 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=(255, 0, 0), thickness=2)
+    
+    canvas = cv2.addWeighted(_image, 1.0, maskCanvas.astype(np.uint8), 1.0, 0)
+    return canvas, maskCanvas
+
+
+def RescaleTheMask(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)
+
+        w = orijinal_image.shape[0]
+        h = orijinal_image.shape[1]
+
+        mask = cv2.drawContours(b_mask, [contour], -1, (1, 1, 1), cv2.FILLED)
+        _masks += [mask]
+
+    return _masks
+
+
+
+image = cv2.imread(test_image)
 
 
 with torch.no_grad():
     results = model(
-        test_image,
-        save=True,
+        image,
+        save=False,
         show_boxes=False,
-        project="./result/",
+        project="./inference/",
         conf=0.5,
+        iou=0.7,
         retina_masks=False
     )
 
+    batches = ParseResults(results, threshold=0.5, scale_masks=True)
+    scores, classes, masks, boxes = batches[0]
 
-with torch.no_grad():
-    for result in results:
-        masks = result.masks.data
-        boxes = result.boxes.data
+    canvas, mask = DrawResults(image, scores, classes, masks, boxes, class_filter=[3])
 
-    #ALL
-    canvas = torch.any(masks, dim=0).int() * 255
 
-    clss = boxes[:, 5]
-    obj_indices = torch.where(clss == 4)
+    # ALL Segmentation
+    # canvas = torch.any(result.masks.data, dim=0).int() * 255
+    
+    # Instance Segmentation
+    # objIdx = torch.where(result.boxes.cls.data == 3)
+    # objMasks = result.masks.data[objIdx]
+    # obj_mask = torch.any(objMasks, dim=0).int() * 255
+    
 
-    # Cable
-    obj_masks = masks[obj_indices]
-    obj_mask = torch.any(obj_masks, dim=0).int() * 255
-    # cropped_image = result.orig_img[obj_mask.cpu().numpy()]
 
-
-fig, axs = plt.subplots(nrows=2, ncols=2, figsize=(27,15))
-axs[0][0].imshow(result.orig_img)
+#! Plot
+fig, axs = plt.subplots(2, 2, figsize=(27, 15))
+axs[0][0].imshow(image)
 axs[0][0].set_title("Orijinal Görüntü")
 
-axs[0][1].imshow(canvas.cpu().numpy())
+axs[0][1].imshow(mask)
 axs[0][1].set_title("Segmentasyon Maskesi")
 
-mask = np.array(obj_mask.cpu().numpy())*255
-cv2.imwrite("cable_mask.png", mask)
-axs[1][0].imshow(obj_mask.cpu().numpy())
-axs[1][0].set_title("Seçilen")
+# axs[1][0].imshow(obj_mask.cpu().numpy())
+# axs[1][0].set_title("Seçilen")
 
-axs[1][1].imshow(result.plot())
+axs[1][1].imshow(canvas)
 axs[1][1].set_title("Sonuç")
+
+# mask = np.array(obj_mask.cpu().numpy())*255
+# cv2.imwrite("cable_mask.png", mask)
+
+plt.tight_layout()
 plt.show()