Adding new Object Detection Models

[AdityaNG]

So far, I have created a basic set of 5 models (YOLOv5 nano through extra-large). We now need to add more such models to the obj_det_demo.py.

To create a new model, following the following structure:

class obj_det_pipeline_model_1(obj_det_evaluator, pipeline_model):
        def load(self):
		# Load your model
		self.model = torch.hub.load('ultralytics/yolov5', 'yolov5s')

	def predict(self, x: dict) -> pd.DataFrame:
                # x: dict whose keys are model name and values are the model prediction
                # returns: pd.DataFrame with the following cols
                predict_results = {
			'xmin': [], 'ymin':[], 'xmax':[], 'ymax':[], 'confidence': [], 'name':[], 'image':[]
		}
                # TODO: Logic
                predict_results = pd.DataFrame(predict_results)
		return predict_results

        def evaluate(self, x, y):
                # Evaluation logic, returns results, press
                # Inherit from obj_det_evaluator if you don't need anything custom
                pass

        def train(self, dataset):
		# TODO: Training, not implemented yet
		pass

Remember to add the newly created class to the obj_det_input = pipeline_input(...) at the very bottom of the obj_det_demo.py

ML_Ops_Pipeline/obj_det_demo.py

Lines 216 to 226 in e29ca6a

	obj_det_input = pipeline_input("obj_det", {'karthika95-pedestrian-detection': obj_det_interp_1},
	{
	'obj_det_pipeline_model_yolov5n': obj_det_pipeline_model_yolov5n,
	'obj_det_pipeline_model_yolov5s': obj_det_pipeline_model_yolov5s,
	'obj_det_pipeline_model_yolov5m': obj_det_pipeline_model_yolov5m,
	'obj_det_pipeline_model_yolov5l': obj_det_pipeline_model_yolov5l,
	'obj_det_pipeline_model_yolov5x': obj_det_pipeline_model_yolov5x,
	}, {'obj_det_pipeline_ensembler_1': obj_det_pipeline_ensembler_1})
	all_inputs = {}
	all_inputs[obj_det_input.get_pipeline_name()] = obj_det_input

My demo implementation of YOLOv5 Model:

ML_Ops_Pipeline/obj_det_demo.py

Lines 128 to 158 in e29ca6a

	class obj_det_pipeline_model(obj_det_evaluator, pipeline_model):
	def load(self):
	self.model = torch.hub.load('ultralytics/yolov5', 'yolov5s')
	def train(self, log_dir, dataset):
	# TODO: Training
	pass
	def predict(self, x: dict) -> np.array:
	# Runs prediction on list of values x of length n
	# Returns a list of values of length n
	predict_results = {
	'xmin': [], 'ymin':[], 'xmax':[], 'ymax':[], 'confidence': [], 'name':[], 'image':[]
	}
	for image_path in x:
	img = cv2.imread(image_path)
	results = self.model(image_path)
	df = results.pandas().xyxyn[0]
	res = df[df["name"]=="person"]
	for index, yolo_bb in res.iterrows():
	file_name = image_path.split('/')[-1][0:-4]
	predict_results["xmin"] += [yolo_bb["xmin"]*img.shape[1]]
	predict_results["ymin"] += [yolo_bb["ymin"]*img.shape[0]]
	predict_results["xmax"] += [yolo_bb["xmax"]*img.shape[1]]
	predict_results["ymax"] += [yolo_bb["ymax"]*img.shape[0]]
	predict_results["confidence"] += [yolo_bb["confidence"]]
	predict_results["name"] += [file_name]
	predict_results["image"] += [image_path]
	predict_results = pd.DataFrame(predict_results)
	return predict_results