# ----------------------------------------------------------------------------- # # This file is part of the PantoScanner distribution on: # https://huggingface.co/spaces/swissrail/PantoScanner # # PantoScanner - Analytics and measurement capability for technical objects. # Copyright (C) 2017-2024 Schweizerische Bundesbahnen SBB # # Authors (C) 2024 L. Hofstetter (lukas.hofstetter@sbb.ch) # Authors (C) 2017 U. Gehrig (urs.gehrig@sbb.ch) # # This program is free software: you can redistribute it and/or modify # it under the terms of the GNU General Public License as published by # the Free Software Foundation, either version 3 of the License, or # (at your option) any later version. # # This program is distributed in the hope that it will be useful, # but WITHOUT ANY WARRANTY; without even the implied warranty of # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the # GNU General Public License for more details. # # You should have received a copy of the GNU General Public License # along with this program. If not, see . # # ----------------------------------------------------------------------------- import os from glob2 import glob import torch import torchvision import cv2 as cv import matplotlib.pyplot as plt import matplotlib import matplotlib.patches as patches import numpy as np from dataclasses import dataclass, field, asdict, astuple, InitVar from typing import List, Dict, Tuple @dataclass class BBoxCoordinates: x_min: int = field(init=True, default=None) x_max: int = field(init=True, default=None) y_min: int = field(init=True, default=None) y_max: int = field(init=True, default=None) x_center: float = field(init=False, default=False) y_center: float = field(init=False, default=False) height: int = field(init=False, default=False) width: int = field(init=False, default=False) as_array: np.ndarray = field(init=False, default=False) def __post_init__(self): self.height = self.x_max - self.x_min self.width = self.y_max - self.y_min self.x_center = 0.5 * (self.x_max + self.x_min) self.y_center = 0.5 * (self.y_max + self.y_min) def shift_coordinates(self, x_shift, y_shift, x_new_low_bound=None, x_new_upp_bound=None, y_new_low_bound=None, y_new_upp_bound=None): self.x_min = self.x_min - x_shift if x_new_low_bound is not None: self.x_min = max(self.x_min, x_new_low_bound) self.x_max = self.x_max - x_shift if x_new_upp_bound is not None: self.x_max = min(self.x_max, x_new_upp_bound) self.y_min = self.y_min - y_shift if y_new_low_bound is not None: self.y_min = max(self.y_min, y_new_low_bound) self.y_max = self.y_max - y_shift if y_new_upp_bound is not None: self.y_max = min(self.y_max, y_new_upp_bound) self.height = self.x_max - self.x_min self.width = self.y_max - self.y_min self.x_center = 0.5 * (self.x_max + self.x_min) self.y_center = 0.5 * (self.y_max + self.y_min) def check_valid(self): if self.width and self.height: return True else: return False def make_yolo_label_string(self, label_number, img_size, float_precision): final_string = str(int(label_number)) for this_float in [self.y_center, self.x_center, self.width, self.height]: final_string = final_string + ' ' + '{:.{n}f}'.format(this_float/img_size, n=float_precision) return final_string def evaluate_yolo(model_path: str, img_path: str, img_size: int, bound_1: Tuple[int, int], bound_2: Tuple[int, int]): model = torch.hub.load('ultralytics/yolov5', 'custom', path=model_path, verbose=None) # local model inpt_img = cv.imread(img_path) img_1 = inpt_img[bound_1[0]:bound_1[0] + img_size, bound_1[1]:bound_1[1] + img_size, :] img_2 = inpt_img[bound_2[0]:bound_2[0] + img_size, bound_2[1]:bound_2[1] + img_size, :] inpt_imgs = [img_1, img_2] results = model(inpt_imgs, size=img_size) result_list_images = [] results_1 = results.pandas().xyxy[0].reset_index() # img1 predictions (pandas) results_2 = results.pandas().xyxy[1].reset_index() # img1 predictions (pandas) for index, row in results_1.iterrows(): result_list_images.append( [row['xmin'] + bound_1[1], row['xmax'] + bound_1[1], row['ymin'] + bound_1[0], row['ymax'] + bound_1[0]]) for index, row in results_2.iterrows(): result_list_images.append( [row['xmin'] + bound_2[1], row['xmax'] + bound_2[1], row['ymin'] + bound_2[0], row['ymax'] + bound_2[0]]) result_arr = np.asarray(result_list_images) x_min = np.min(result_arr[:, 0]) x_max = np.max(result_arr[:, 1]) y_min = np.min(result_arr[:, 2]) y_max = np.max(result_arr[:, 3]) bound_final_x = max(min(int(0.5 * (x_min + x_max - img_size)), 3208 - img_size), 0) bound_final_y = max(min(int(0.5 * (y_min + y_max - img_size)), 2200 - img_size), 0) img_final = inpt_img[bound_final_y:bound_final_y + img_size, bound_final_x:bound_final_x + img_size, :] imgs_final = [img_final] results_final = model(imgs_final, size=img_size) bboxes_final = results_final.pandas().xyxy[0].reset_index() # img1 predictions (pandas) final_bbox_list = [] for index, row in bboxes_final.iterrows(): this_box = BBoxCoordinates(row['ymin'] + bound_final_y, row['ymax'] + bound_final_y, row['xmin'] + bound_final_x, row['xmax'] + bound_final_x) final_bbox_list.append(this_box) return final_bbox_list def evaluate_yolo_2(model_loaded, img_path: str, img_size: int, bound_1: Tuple[int, int], bound_2: Tuple[int, int]): model = model_loaded inpt_img = cv.imread(img_path) img_1 = inpt_img[bound_1[0]:bound_1[0] + img_size, bound_1[1]:bound_1[1] + img_size, :] img_2 = inpt_img[bound_2[0]:bound_2[0] + img_size, bound_2[1]:bound_2[1] + img_size, :] inpt_imgs = [img_1, img_2] results = model(inpt_imgs, size=img_size) result_list_images = [] results_1 = results.pandas().xyxy[0].reset_index() # img1 predictions (pandas) results_2 = results.pandas().xyxy[1].reset_index() # img1 predictions (pandas) for index, row in results_1.iterrows(): result_list_images.append( [row['xmin'] + bound_1[1], row['xmax'] + bound_1[1], row['ymin'] + bound_1[0], row['ymax'] + bound_1[0]]) for index, row in results_2.iterrows(): result_list_images.append( [row['xmin'] + bound_2[1], row['xmax'] + bound_2[1], row['ymin'] + bound_2[0], row['ymax'] + bound_2[0]]) result_arr = np.asarray(result_list_images) x_min = np.min(result_arr[:, 0]) x_max = np.max(result_arr[:, 1]) y_min = np.min(result_arr[:, 2]) y_max = np.max(result_arr[:, 3]) bound_final_x = max(min(int(0.5 * (x_min + x_max - img_size)), 3208 - img_size), 0) bound_final_y = max(min(int(0.5 * (y_min + y_max - img_size)), 2200 - img_size), 0) img_final = inpt_img[bound_final_y:bound_final_y + img_size, bound_final_x:bound_final_x + img_size, :] imgs_final = [img_final] results_final = model(imgs_final, size=img_size) bboxes_final = results_final.pandas().xyxy[0].reset_index() # img1 predictions (pandas) final_bbox_list = [] for index, row in bboxes_final.iterrows(): this_box = BBoxCoordinates(row['ymin'] + bound_final_y, row['ymax'] + bound_final_y, row['xmin'] + bound_final_x, row['xmax'] + bound_final_x) final_bbox_list.append(this_box) return final_bbox_list # # # img_folder_path = 'C:/Users/hofst/PycharmProjects/ImageAnalysis_SBB/training_data/ext_label_1/images/' # img_list = [os.path.normpath(this_path) for this_path in glob(img_folder_path + '*.png')] # b_box_list = [] # model_path_1 = 'C:\\Users\\hofst\\PycharmProjects\\ImageAnalysis_SBB\\yolov5\\data\\richard\\yolo_training_test_1408_size_m\\weights\\best.pt' # img_path_new = 'C:\\Users\\hofst\\PycharmProjects\\ImageAnalysis_SBB\\training_data\\test_1.png' # bound_1_new = (300, 92) # bound_2_new = (650, 1500) # img_size_new = 1408 # # this_model = torch.hub.load('ultralytics/yolov5', 'custom', path=model_path_1, verbose=False) # local model # # for img_path_new in img_list: # test_result = evaluate_yolo_2(this_model, img_path_new, img_size_new, bound_1_new, bound_2_new) # b_box_list.append(test_result) # matplotlib.use('TkAgg') # Change backend after loading model # fig, ax = plt.subplots() # # Display the image # ax.imshow(cv.imread(img_path_new), cmap='gray') # # Create a Rectangle patch # rect = patches.Rectangle((test_result[0].y_min, test_result[0].x_min), test_result[0].y_max - test_result[0].y_min, test_result[0].x_max - test_result[0].x_min, linewidth=1, edgecolor='r', facecolor='none') # rect_2 = patches.Rectangle((test_result[1].y_min, test_result[1].x_min), test_result[1].y_max - test_result[1].y_min, test_result[1].x_max - test_result[1].x_min, linewidth=1, edgecolor='b', facecolor='none') # # # Add the patch to the Axes # ax.add_patch(rect) # ax.add_patch(rect_2) # fig.show() # # #x_1_list = [] # #y_1_list = [] # #x_2_list = [] # #y_2_list = [] # # #for box_pair in b_box_list: # # box_1 = box_pair[0] # # box_2 = box_pair[1] # # x_1_list.append(box_1.y_min) # # y_1_list.append(box_1.x_max) # # x_1_list.append(box_2.y_min) # # y_1_list.append(box_2.x_max) # # # x_2_list.append(box_1.y_max) # # y_2_list.append(box_1.x_min) # # x_2_list.append(box_2.y_max) # # y_2_list.append(box_2.x_min) # # #ax.scatter(x_1_list, y_1_list) # #ax.scatter(x_2_list, y_2_list) # #plt.imshow(cv.imread(img_path_new), cmap='gray') # #plt.grid()