Spaces:
Running
Running
| 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 |