handler.py

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)
        
        url = "https://drive.google.com/file/d/1jB8sDYYOTfuF7B1PMcDjkm5R7huv97Wm"
        gdown.download(url, 'best.pt', quiet=False)
        
        self.model = YOLO("best.pt")
    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
        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))