src/apps/detection.py

'''
@Author     : Ali Mustofa HALOTEC
@Module     : Container Number Iso Code Detection
@Created on : 10 Oct 2022
'''
#!/usr/bin/env python3
# Path: src/apps/detection.py
import os
import cv2
import numpy as np
from PIL import Image
from typing import Optional
from pydantic import BaseModel
from src.utils.utils import download_and_unzip_model

import torch
from torch import nn

class InputDetection(BaseModel):
	image: 			list
	min_confidence: Optional[float] = 0.5
	get_one: 		Optional[bool] 	= True
	boxes_ori: 		Optional[bool] 	= True
	resized_size: 	Optional[int] 	= 640
	image_shape:	Optional[tuple] = (1280, 720, 3)
	
class Detection:
	def __init__(self, root_path:str, model_config:dict) -> None:
		'''
		Load model 
		@params:
			- root_path:str ->  root of path model
			- model_config:dict -> config of model {filename, classes, url, file_size}
		'''
		self.root_path :str          = root_path
		self.model_config :dict      = model_config
		self.model_name :str         = f'{root_path}/{model_config["filename"]}'
		self.image_size :int         = model_config['image_size']
		self.classes :list           = model_config['classes']
		self.filter_classes :list    = model_config['filter_classes']
		self.device :torch           = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
		self.model :nn.Module        = self.__load_model()

	@staticmethod
	def __check_model(root_path:str, model_config:dict) -> None:
		if not os.path.isfile(f'{root_path}/{model_config["filename"]}'):
			download_and_unzip_model(
				root_dir    = root_path,
				name        = model_config['filename'],
				url         = model_config['url'],
				file_size   = model_config['file_size'],
				unzip       = False
			)
		else: print('Load model char detection')

	def __load_model(self) -> torch.nn.Module:
		self.__check_model(self.root_path, self.model_config)
		try:
			model = torch.hub.load('ultralytics/yolov5', 'custom', path=self.model_name)
		except:
			model = torch.hub.load('ultralytics/yolov5', 'custom', path=self.model_name, force_reload=True)
		return model
	
	def extract_results(self, results, min_confidence:float = 0.5, get_one=False, \
		boxes_ori:bool = True, resized_size:int = 0, image_shape:tuple = (1080, 720)):
		'''
		Format result([tensor([[151.13147, 407.76913, 245.91382, 454.27802,   0.89075,   0.00000]])])
		Filter min confidence prediction and classes id/name
		Cropped image and get index max value confidence lavel
		Args:
			result(models.common.Detections): result detection YoloV5
			min_confidence(float): minimal confidence detection in range 0-1
			boxes_ori(bool): if true, calculate boxes to original resolution
			resized_size(int): value of resized image detection
			image_shape(tuple): height, width image original
		Return:
			result(dict): {
				casess:[{
					confidence(float): confidence,
					bbox(list) : [x_min, y_min, x_max, y_max]
				}]
			}
		'''
		results_format  = results.xyxy
		results_filter =  dict({i:list() for i in self.classes})
		
		if len(results_format[0]) >= 1:
			for i in range(len(results_format[0])):
				classes_name    = self.classes[int(results_format[0][i][-1])]
				confidence      = float(results_format[0][i][-2])
				if classes_name in self.filter_classes and confidence >= min_confidence:

					x_min, y_min = int(results_format[0][i][0]), int(results_format[0][i][1])
					x_max, y_max = int(results_format[0][i][2]), int(results_format[0][i][3])
					# change coordinate to original w n h
					if boxes_ori:
						resized_size = self.image_size if resized_size == 0 else resized_size
						height, width = image_shape[:2]
						x_min = int((x_min/resized_size)*width)
						y_min = int((y_min/resized_size)*height)
						x_max = int((x_max/resized_size)*width)
						y_max = int((y_max/resized_size)*height)

					if get_one:
						if results_filter[classes_name]:
							if results_filter[classes_name][0]['confidence'] < confidence:
								results_filter[classes_name][0] = \
									{'confidence': round(confidence,2), 
									'bbox':[x_min, y_min, x_max, y_max]}
							# else: pass
						else:
							results_filter[classes_name].append({
								'confidence': round(confidence,2), 
								'bbox':[x_min, y_min, x_max, y_max]
							})
					else:
						results_filter[classes_name].append({
							'confidence': round(confidence,2), 
							'bbox':[x_min, y_min, x_max, y_max]
						})
		# Delete key if detection null
		for i in self.classes:results_filter.__delitem__(i) if not results_filter[i] else None

		return results_filter
	
	@staticmethod
	def visualize_result(image:np.array, results_filter:dict) -> np.array:
		'''
		Draw bounding box result
		Args:
			image(numpy.ndarray) : image/frame
			results_filter(dict) : {
				casess:[{
					confidence(float): confidence,
					bbox(list) : [x_min, y_min, x_max, y_max]
				}]
			}
		Return:
			image(numpy.ndarray) : image drawed
		'''
		for result in results_filter.items():
			label = result[0]
			for i in result[1]:
				conf = int(i['confidence']*100)
				bbox = i['bbox']
				x_min, y_min, x_max, y_max = \
				bbox[0], bbox[1], bbox[2], bbox[3]

				cv2.rectangle(image, (x_min, y_min), (x_max, y_max), (86, 71, 255), 2)
				# Draw text
				label_conf = f'{label.upper()}:{conf}%'
				cv2.rectangle(image, (x_min, y_min-20), (x_min+(len(label_conf)*13), y_min), (86, 71, 255), cv2.FILLED)
				cv2.putText(image, label_conf, (x_min, y_min), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 255), 2)
		return image

	def __call__(self, image:np.array, size:int = None):
		'''
		Prediction image object detectionn YoloV5
		Args:
			image(numpy.ndarray) : image/frame
		Return:
			results_prediction(models.common.Detections) : results -> convert to (results.xyxy/resultsxywh)
		'''
		if size: results = self.model(image, size=size)
		else: results = self.model(image)
		return results

if __name__ == '__main__':
	root_path = os.path.expanduser('/Users/alimustofa/Downloads/')
	detection_config = {
		'filename'          : 'avori_detection.pt',
		'image_size'        : 640,
		'classes'           : ['avocado'],
		'filter_classes'    : ['avocado'],
		'url'               : 'https://github.com/Alimustoofaa/ContainerNumber-Dev/releases/download/detection_v1/container_iso_maxgross.pt',
		'file_size'         : 14749585,
		
	}
	detection = Detection(root_path, detection_config)

	image = cv2.imread('12022051823052897.jpg')
	image_resize = cv2.resize(
					image.copy(),
					(detection_config['image_size'], 
					detection_config['image_size']), 
					interpolation = cv2.INTER_AREA)

	results = detection(image_resize, size=detection_config['image_size'])
	filtered = detection.extract_results(
					results, 0.5,
					get_one=True, 
					boxes_ori=True, 
					resized_size=detection_config['image_size'], 
					image_shape=image.shape)
	img_drawed = detection.visualize_result(image.copy(), filtered)
	cv2.imwrite('img_drawed.jpg', img_drawed)
	Image.fromarray(cv2.cvtColor(img_drawed, cv2.COLOR_BGR2RGB))