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 | |
from PIL import Image | |
import io | |
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] | |
clothes = cv2.cvtColor(clothes, cv2.COLOR_BGR2RGB) | |
retval , buffer = cv2.imencode('.jpg', clothes) | |
im_bytes = buffer.tobytes() | |
# cv2.imwrite("result.jpg", clothes) | |
# create base 64 object | |
jpg_as_text = base64.b64encode(im_bytes).decode("utf-8") # Decode bytes to string") | |
# base64_encoded = base64.b64encode(image_bytes).decode("utf-8") | |
# Get the image format | |
image_format = Image.open(io.BytesIO(buffer)).format.lower() | |
# Construct the data URI | |
data_uri = f"data:image/{image_format};base64,{jpg_as_text}" | |
# return data_uri | |
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)) | |