Spaces:
Sleeping
Sleeping
File size: 3,593 Bytes
82e13f0 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 |
import os, shutil, random
from PIL import Image
from ultralytics import YOLO
from ultralytics.utils.plotting import Annotator
class SaigonBeer_Recognition:
def __init__(self, cfg):
self.cfg = cfg
self.classifier = YOLO(self.cfg["Model"]['classifier'])
self.detector = YOLO(self.cfg["Model"]['detector'])
self.color = {#"others": (255,0,0),
"saigon_export":(0,255,0),
"saigon_chill":(0,0,255),
"saigon_large":(255,255,0),
"saigon_special":(255,0,255),
"saigon_gold":(255,0,0)}
def forward(self, img, save_path=None):
detected_objs = self.detector.predict(img)
object_info = self.get_objects(detected_objs)
object_info = self.box_filtering(object_info) # dictionary
result = {}
annotated_image = img.copy()
for key, value in object_info.items():
cropped_img = self.image_cropping(img, value['bbox'])
kind = value['label']
beer = self.classifier.predict(source=cropped_img, conf=0.5)
beer = self.got_final_result(beer)
result[key] = {"shape": kind, "beer": beer}
if beer == "others":
continue
annotated_image = self.draw_img(annotated_image, value['bbox'], kind, beer, color=self.color[beer])
return result, annotated_image
def got_final_result(self, results):
'''
Get the label of the largest probability.
Args:
results: The results of the prediction.
Returns:
The label of the largest probability.
'''
id2label = results[0].names
for result in results:
cls = result.probs.top1
return id2label[cls]
def image_cropping(self, img, bbox):
'''
Crop the image based on the bounding box.
Args:
img: The path of the image.
bbox: The bounding box of the object.
Returns:
The cropped image.
'''
x_min, y_min, x_max, y_max = bbox
cropped_img = img.crop((x_min, y_min, x_max, y_max))
return cropped_img
def get_objects(self, predictions):
'''
Get the information of detected object.
Args:
predictions: The results of the prediction.
Returns:
A dictionary containing the information of detected object.
'''
object_dict = {}
for i, prediction in enumerate(predictions):
boxes = prediction.boxes
for j, box in enumerate(boxes):
object_key = f"object_{i}_{j}"
object_dict[object_key] = {
'label': prediction.names[int(box.cls)],
'bbox': box.xyxy.tolist()[0],
'prob': box.conf.tolist()[0]
}
return object_dict
def box_filtering(self, predictions):
'''
Filtering out the lowest boxes
Args:
predictions: The results of the prediction.
Returns:
The filtered predictions.
'''
# Create a copy of the keys to iterate over
keys_to_remove = []
for key, value in predictions.items():
if value['prob'] < 0.5:
keys_to_remove.append(key)
# Remove keys outside of the loop
for key in keys_to_remove:
del predictions[key]
return predictions
def draw_img(self, img, bbox, shape, beer, color=(0,255,0)):
self.annotator = Annotator(img)
self.annotator.box_label(bbox, f"{beer} {shape}", color=color)
annotated_image = self.annotator.result()
# Convert the NumPy array to a PIL Image before saving
annotated_image = Image.fromarray(annotated_image)
return annotated_image |