added auto annotation logic

inference.py

@@ -14,7 +14,7 @@ from utils.general import (check_img_size, non_max_suppression)
 from utils.torch_utils import select_device, smart_inference_mode
 
 class ExecuteInference:
-    def __init__(self, weight, confidence=0.4, img_size=640, agnostic_nms=False, gpu=False, iou=0.5):
+    def __init__(self, weight, confidence=0.4, img_size=416, agnostic_nms=False, gpu=False, iou=0.5):
         self.weight = weight
         self.confidence = confidence
         self.gpu = gpu
@@ -45,6 +45,7 @@ class ExecuteInference:
         return model, names, colors
 
     def predict(self, image):
+        image2 = image
         img = letterbox(image, new_shape=self.img_size)[0]
         img = img[:, :, ::-1].transpose(2, 0, 1)
         img = np.ascontiguousarray(img)
@@ -58,14 +59,82 @@ class ExecuteInference:
         _output = list()
         for i, det in enumerate(pred):
             if det is not None and len(det):
-                det[:, :4] = scale_boxes(img.shape[2:], det[:, :4], img.shape).round()
+                print(det)
+                det[:, :4] = scale_boxes(img.shape[2:], det[:, :4], image2.shape).round()
                 for *xyxy, conf, cls in reversed(det):
+                    print(xyxy)
                     _output.append({"points": xyxy, "conf": conf, "class": cls})
         return _output
 
+import os
 
-x = ExecuteInference(weight="yolov5s.pt")
-img = cv2.imread(r"C:\Users\sikhin.vc\PycharmProjects\cuda_enabled_repo\custom_image_augmentation\yolov5l_results12\confusion_matrix.png")
-output = x.predict(img)
 
-print(output)
+def get_scaled_point(iheight, iwidth, points):
\ No newline at end of file
+    x1, y1, x2, y2 = points
+    x1 = x1 / iwidth
+    x2 = x2 / iwidth
+    y1 = y1 / iheight
+    y2 = y2 / iheight
+
+    centerx = (x2 + x1) / 2
+    centery = (y2 + y1) / 2
+
+    widthi = y2 - y1
+    heighti = x2 - x1
+    return centerx, centery, widthi, heighti
+
+
+
+
+yp = ExecuteInference(weight=r"C:\Users\sikhin.vc\PycharmProjects\yolov5_auto_annotation\voc_to_yolo\packer_5.pt")
+directory = r"C:\Users\sikhin.vc\Pictures\Packer5_Merged\Packer5_Merged"
+target_dir = r"C:\Users\sikhin.vc\Pictures\Packer5_Merged\Packer5_Merged"
+# img = cv2.imread(r"C:\Users\sikhin.vc\Documents\loader7_vinod_frames\loader7_vinod_frames\images\Frame_vinod_3768.jpg")
+# output = x.predict(img)
+
+# print(output)
+
+for each_file in os.listdir(directory):
+    title, ext = os.path.splitext(os.path.basename(each_file))
+    if ext != ".jpg":
+        continue
+    print(each_file)
+    image = cv2.imread(os.path.join(directory, each_file))
+    image = cv2.resize(image, (416, 416))
+    img = image.copy()
+    predict = yp.predict(image)
+    print(predict)
+    c = []
+    for p in predict:
+        c.append([int(p["points"][0]), int(p["points"][1]), int(p["points"][2]), int(p["points"][3])])
+    # print(predict)
+    print(c)
+    height, width, _ = image.shape
+    # cv2.imwrite(os.path.join(target_dir, title + ".jpg"), image)
+    # for each_pred1 in predict:
+        # if(each_pred1["class"] in cls_li):
+
+    with open(os.path.join(target_dir, title + ".txt"), 'a') as f:
+        for each_pred in c:
+            mrp_label = []
+            # labels =  ["cell phone"]
+            points = [each_pred[0], each_pred[1], each_pred[2], each_pred[3]]
+            # img = cv2.rectangle(img, (each_pred[0], each_pred[1]),
+            #                     (each_pred[2], each_pred[3]), (255, 0, 0), 2)
+            # cv2.imshow("out", img)
+            # cv2.waitKey(0)
+            cx, cy, w, h = get_scaled_point(height, width, each_pred)
+            # if each_pred["class"] in cls_li:
+            f.write("0 {} {} {} {}\n".format(cx, cy, h, w))
+            # else:
+            #
+            #     mrp_label.append([cx, cy, w, h])
+
+                # else:
+                #
+            #     #     f.write("0 {} {} {} {}\n".format(cx, cy, h, w))
+            # for label in mrp_label:
+            #     f.write("1 {} {} {} {}\n".format(label[0], label[1], label[3], label[2]))
+
+        cv2.imshow("image", img)
+        cv2.waitKey(1)
\ No newline at end of file