from typing import Dict, List, Any import urllib.request import numpy as np import cv2 import base64 from ultralytics import YOLO import os import gdown class EndpointHandler: def __init__(self, path='.'): # pass api key to model # current_directory = os.getcwd() # print("Current working directory:", current_directory) url = "https://drive.google.com/uc?id=1jB8sDYYOTfuF7B1PMcDjkm5R7huv97Wm" gdown.download(url, './best.pt', quiet=False) self.model = YOLO("./best.pt") def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]: os.getcwd() 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.shape[0] == 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 = EndpointHandler() data = { "inputs": { "isurl": True, # "path": "http://10.10.2.100/cam-lo.jpg", "path": "https://www.next.us/nxtcms/resource/blob/5791586/ee0fc6a294be647924fa5f5e7e3df8e9/hoodies-data.jpg", # "key": "iJuYzEzNEFSaQq4e0hfE", } } # test file image print(Model(data)) #test url # print(Model("http://10.10.2.100/cam-lo.jpg", 1))