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/view?usp=sharing"
        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))