Image-preprocessor / handler.py
oscarfu0501's picture
Update handler.py
20fa134 verified
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
# import http.client
# http.client.HTTPConnection._http_vsn = 10
# http.client.HTTPConnection._http_vsn_str = 'HTTP/1.0'
class EndpointHandler:
def __init__(self, path='.'): # pass api key to model
# current_directory = os.getcwd()
# print("Current working directory:", current_directory)
pass
def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
url = "https://drive.google.com/file/d/1jB8sDYYOTfuF7B1PMcDjkm5R7huv97Wm/view?usp=sharing"
gdown.download(url, './best.pt', quiet=False)
model = YOLO("./best.pt")
inputs = data.get("inputs")
print("in call")
isurl = inputs.get("isurl")
print("in isurl")
path = inputs.get("path")
print("is path")
print(path)
# path = "http://10.10.2.100/cam-lo.jpg"
# 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(buffer).decode("utf-8") # Decode bytes to string")
jpg_as_text = base64.b64encode(buffer).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))