Commited plain concepts trial

data_analysis.py

@@ -74,43 +74,10 @@ class DataAnalyst():
             except Exception as e:
                 print(f"Error while trying to resize and save and image {e}")
 
-    def xml_data_extractor(self):
-        """THIS IS THE FUNCTION THAT WILL EXTRACT LABELS DATA FROM THE XML FILES
-           THESE XML FILES DESCRIBE THE BOUNDING BOXES AROUND THE CIRCUIT ELEMENTS"""
-        for idx,file in enumerate(os.listdir(self.annotations_route)):
-            file_labels = []
-            file_path = os.path.join(self.annotations_route, file)
-            try:
-                with open(file_path, 'r') as f:
-                    file_data = f.read()
-                    beauti_data = BeautifulSoup(file_data, 'lxml-xml')
-                    elements_list = beauti_data.find_all('object')
-                
-                    for element in elements_list:
-                        name = element.find("name").text
-                        if name =="Capactitor":
-                            name ="Capacitor"
-                        elif name == "DC voltage source":
-                            self.dc_circuits.append(idx)
-                        elif name == "AC voltage source":
-                            self.ac_circuits.append(idx)
-
-                        self.elements_names.append(name)
-                        coord_x1 = element.find("xmin").text
-                        coord_x2 = element.find("xmax").text
-                        coord_y1 = element.find("ymin").text
-                        coord_y2 = element.find("ymax").text
-                        file_labels.append([float(coord_x1), float(coord_x2), float(coord_y1), float(coord_y2), float(self.names_dict[name])])
-                self.training_labels.append(file_labels)
-            except Exception as e:
-                print(f"Error while reading xml files and getting their information {e}")
-
-
     def dataset_info_summary(self):
         """HERE I AM GOING TO VISUALIZE SOME DATASET INFORMATION, TYPICAL HEIGHT AN WIDTH, CLASS DISTRIBUTION"""
         elements_dataframe = pd.DataFrame(self.elements_names)
         elements_names =elements_dataframe.value_counts()
-        print(elements_names)
         fig,axes = plt.subplots(1,1)
         axes.pie(elements_names, labels= elements_names.index, autopct='%1.2f%%', colors=['gold', 'skyblue', 'lightgreen', 'red', 'green'], startangle=90)
         plt.show()
@@ -142,8 +109,6 @@ class DataAnalyst():
                     f.close()
 
 
-
-
     def xml_data_getter(self, file_path):
         """THIS IS THE FUNCTION THAT WILL EXTRACT LABELS DATA FROM THE XML FILES
            THESE XML FILES DESCRIBE THE BOUNDING BOXES AROUND THE CIRCUIT ELEMENTS"""
@@ -175,11 +140,8 @@ class DataAnalyst():
         """THIS FUNCTION WILL PREPARE THE COORDINATES THEY WAY YOLO EXPECTS IT
         FIRST GET COORDINATES NORMALIZED, SINCE YOLO EXPECTS TO BE BETWEEN 0 AND 1
         THIS IS BECAUSE YOLO RESEARCHES REALISED THAT IT WAYS EASIER TO USE OFFSETS OF GRID CELLS RATHER THAN PURE COORDINATES"""
-        print(bbox_tensor)
         normalized_width = 1.0 / image_size[0]
         normalized_height = 1.0 / image_size[1]
-        print(normalized_width)
-        print(normalized_height)
         width = (bbox_tensor[1] - bbox_tensor[0]) * normalized_width
         height = (bbox_tensor[3] - bbox_tensor[2]) * normalized_height
 

dataprepare.py

@@ -1,3 +1,7 @@
+"""THIS FILE IS PART OF THE FIRST APPROXIMATION THAT I DID TO THE PROBLEM USING DATASET CLASS AND ITERATING OVER IT USING A DATALOADER
+   BUT  IT IS  NOT USED IN THE FINAL SOLUTION AFTER REALISING THAT IT WOULDN'T WORK IN ORDER TO FINE TUNE YOLOV8 FROM ULTRALYTICS
+   I KEEP IT HERE SO YOU CAN SEE HOW I USE THIS CLASSES"""
+
 from torch.utils.data import DataLoader, Dataset
 from torchvision import transforms
 from data_analysis import DataAnalyst
@@ -73,21 +77,6 @@ def data_asociation(mode):
             shutil.copy(new_file, labels)
 
 
-
-
-
-
-def analyst_starting(route, output_folder, annotations_route):
-    #THIS FUNCTION IS GOING TO BE CALLED BY THE MAIN IN ORDER TO START THE JOB
-    data_analyst = DataAnalyst(route, output_folder, annotations_route)
-    #data_analyst.image_size_searching()
-    #data_analyst.xml_data_extractor()
-    #data_analyst.dataset_info_summary()
-    #dataset = CircuitDataset(data_analyst.training_features, data_analyst.training_labels, data_analyst.route,(384,384))
-    #training_loader = DataLoader(dataset, batch_size=2, shuffle=True)
-    #for x,y in training_loader:
-    #    print("Imprimo imagen")
-
 def analyst_starting2(route,output_folder,annotations_route):
     data_analyst = DataAnalyst(route, output_folder, annotations_route)
 

train.py

@@ -1,22 +1,65 @@
 from ultralytics import YOLO, settings
 import ultralytics
-from pathlib import Path
-import os
 import torch
+from sys import ps1
+#DATA AUGMENTATION SECTION
+from ultralytics.data.augment import Albumentations
+from ultralytics.utils import LOGGER,colorstr
+import albumentations as A
+import cv2
+
+def create_transformations(p=1.0):
+  """THIS IS WHERE I CREATED THE FIRST TRASNFORMATIONS OVER THE IMAGE, FINALLY I AM JUST USING THE ULTRALYTICS DEFAULT ONE"""
+  prefix = colorstr("albumentations: ")
+  try:
+    transforms = [
+      A.RandomRain(p=0.1, slant_lower=-10, slant_upper=10,
+                              drop_length=20, drop_width=1, drop_color=(200, 200, 200),
+                              blur_value=5, brightness_coefficient=0.9, rain_type="default"),
+      A.Rotate(limit = 10, p=0.5),
+      A.Blur(p=0.1),
+      A.MedianBlur(p=0.1),
+      A.ToGray(p=0.01),
+      A.CLAHE(p=0.01),
+      A.ImageCompression(quality_lower=75, p=0.0),
+    ]
+    transform = A.Compose(transforms, bbox_params=A.BboxParams(format="yolo", label_fields=["class_labels"]))
+    return transform
+
+  except Exception as e:
+    print(f"Error while creating transformations {e}")
 
 def model_training():
+    """FUNCTION USED TO TRAIN THE MODEL"""
     try:
-        #ultralytics.settings.update({'datasets_dir': "./dataset"})
-        #save_path = '/tmp/runs/detect/train2'
-        #Path(save_path).mkdir(parents=True, exist_ok=True)
+        transformations = create_transformations()
         device = "cuda" if torch.cuda.is_available() else "cpu"
         print("Available device: ", device)
         model = YOLO("yolov8n.pt")
-        model.train(data="datasets/dataset.yaml", epochs=20, imgsz=244, batch = 2, device=device,verbose=True)
+        results = model.train(data="/content/datasets/dataset.yaml", epochs=50, imgsz=384, batch = 2, augment = True,device=device,verbose=True)
 
+        return model
     except Exception as e:
-        print(f"Error when training {e}")
+        print(e)
+
+def image_show_fn(bboxes,path):
+    """THIS IS THE FUNCTION USING BY model_testing FUNCTION TO DRAW BBOXES"""
+    image = cv2.imread(path)
+    for bbox in bboxes:
+        bbox = [int(num) for num in bbox]
+        cv2.rectangle(image, (bbox[0], bbox[1]), (bbox[2], bbox[3]), (0, 255, 0), 2)
+    cv2.imshow(image)
+    #THIS METHOD IS ONLY CALLABLE IN GOOGLE COLAB TYPES
+    #cv2_imshow(image)
 
+def model_testing(model, test_image):
+    """THIS IS THE FUNCTION THAT THE MODEL WILL USE AT INFERENCE TIME
+        IT LOADS THE BEST WEIGHTS, DO THE INFERENCE ON THE IMAGE, CALL THE
+        FUNCTION THAT WILL DRAW THE BOUNDING BOXES AND FINALLY DRAW BBOXES """
+    model = YOLO("/content/runs/detect/train12/weights/best.pt")
+    image_results = model(test_image)
+    for image in image_results:
+        image_show_fn(image.boxes.xyxy,image.path)