| from typing import Dict, List, Any | |
| import urllib.request | |
| import numpy as np | |
| import cv2 | |
| import base64 | |
| from ultralytics import YOLO | |
| import os | |
| class EndpointHandler: | |
| def __init__(self, path='.'): # pass api key to model | |
| current_directory = os.getcwd() | |
| print("Current working directory:", current_directory) | |
| self.model = YOLO("../best.pt") | |
| def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]: | |
| inputs = data.get("inputs") | |
| isurl = inputs.get("isurl") | |
| path = inputs.get("path") | |
| model = self.model | |
| ########################### Load Image ################################# | |
| if(isurl): # for url set isurl = 1 | |
| print("checkpoint 2-1") | |
| req = urllib.request.urlopen(path) | |
| print("checkpoint 2-2") | |
| arr = np.asarray(bytearray(req.read()), dtype=np.uint8) | |
| print("checkpoint 2-3") | |
| img = cv2.imdecode(arr, -1) # 'Load it as it is' | |
| else: # for image file | |
| img = cv2.imread(path) | |
| print("checkpoint 2") | |
| ########################################################################### | |
| ########################### Model Detection ################################# | |
| # change model_id to use a different model | |
| # can try: | |
| # clothing-detection-s4ioc/6 //good | |
| # clothing-segmentation-dataset/1 | |
| # t-shirts-detector/1 | |
| # mainmodel/2 | |
| #result = self.CLIENT.infer(path, model_id="mainmodel/2") | |
| result = model(img) | |
| #annotated_frame = result[0].plot() | |
| detections = result[0].boxes | |
| #print(result[0].boxes.xyxy) | |
| #cv2.imshow("YOLOv8 Inference", annotated_frame) | |
| # print(result) | |
| #cv2.waitKey(0) | |
| #detections = sv.Detections.from_inference(result) | |
| # print(detections) | |
| print("checkpoint 3") | |
| ########################################################################### | |
| ########################### Data proccessing ################################# | |
| # only pass the first detection | |
| # change 1 -> to len(detections.xyxy) to pass all photos | |
| if(detections.xyxy.size == 0): | |
| return "Not Found" | |
| else: | |
| x1, y1, x2, y2 = int(detections.xyxy[0][0]), int(detections.xyxy[0][1]), int(detections.xyxy[0][2]), int(detections.xyxy[0][3]) | |
| clothes = img[y1: y2, x1: x2] | |
| retval , buffer = cv2.imencode('.jpg', clothes) | |
| cv2.imwrite("result.jpg", clothes) | |
| # create base 64 object | |
| jpg_as_text = base64.b64encode(buffer) | |
| print("checkpoint 4") | |
| ########################################################################### | |
| return jpg_as_text | |
| ########################################################################### | |
| # test run | |
| # Model = Image_detect() | |
| # test file image | |
| # print(Model("test_images/test6.jpg", 0)) | |
| #test url | |
| # print(Model("http://10.10.2.100/cam-lo.jpg", 1)) | |