test code

.vscode/PythonImportHelper-v2-Completion.json

@@ -0,0 +1,554 @@
+[
+    {
+        "label": "GoogleGenerativeAI",
+        "importPath": "langchain_google_genai",
+        "description": "langchain_google_genai",
+        "isExtraImport": true,
+        "detail": "langchain_google_genai",
+        "documentation": {}
+    },
+    {
+        "label": "onnxruntime",
+        "kind": 6,
+        "isExtraImport": true,
+        "importPath": "onnxruntime",
+        "description": "onnxruntime",
+        "detail": "onnxruntime",
+        "documentation": {}
+    },
+    {
+        "label": "preprocess",
+        "importPath": "src.utils.utils_segment",
+        "description": "src.utils.utils_segment",
+        "isExtraImport": true,
+        "detail": "src.utils.utils_segment",
+        "documentation": {}
+    },
+    {
+        "label": "postprocess",
+        "importPath": "src.utils.utils_segment",
+        "description": "src.utils.utils_segment",
+        "isExtraImport": true,
+        "detail": "src.utils.utils_segment",
+        "documentation": {}
+    },
+    {
+        "label": "extract_text",
+        "importPath": "src.utils.utils_segment",
+        "description": "src.utils.utils_segment",
+        "isExtraImport": true,
+        "detail": "src.utils.utils_segment",
+        "documentation": {}
+    },
+    {
+        "label": "draw_bounding_boxes",
+        "importPath": "src.utils.utils_segment",
+        "description": "src.utils.utils_segment",
+        "isExtraImport": true,
+        "detail": "src.utils.utils_segment",
+        "documentation": {}
+    },
+    {
+        "label": "numpy",
+        "kind": 6,
+        "isExtraImport": true,
+        "importPath": "numpy",
+        "description": "numpy",
+        "detail": "numpy",
+        "documentation": {}
+    },
+    {
+        "label": "Image",
+        "importPath": "PIL",
+        "description": "PIL",
+        "isExtraImport": true,
+        "detail": "PIL",
+        "documentation": {}
+    },
+    {
+        "label": "cv2",
+        "kind": 6,
+        "isExtraImport": true,
+        "importPath": "cv2",
+        "description": "cv2",
+        "detail": "cv2",
+        "documentation": {}
+    },
+    {
+        "label": "Tuple",
+        "importPath": "typing",
+        "description": "typing",
+        "isExtraImport": true,
+        "detail": "typing",
+        "documentation": {}
+    },
+    {
+        "label": "pytesseract",
+        "importPath": "pytesseract",
+        "description": "pytesseract",
+        "isExtraImport": true,
+        "detail": "pytesseract",
+        "documentation": {}
+    },
+    {
+        "label": "load_dotenv",
+        "importPath": "dotenv",
+        "description": "dotenv",
+        "isExtraImport": true,
+        "detail": "dotenv",
+        "documentation": {}
+    },
+    {
+        "label": "FastAPI",
+        "importPath": "fastapi",
+        "description": "fastapi",
+        "isExtraImport": true,
+        "detail": "fastapi",
+        "documentation": {}
+    },
+    {
+        "label": "UploadFile",
+        "importPath": "fastapi",
+        "description": "fastapi",
+        "isExtraImport": true,
+        "detail": "fastapi",
+        "documentation": {}
+    },
+    {
+        "label": "status",
+        "importPath": "fastapi",
+        "description": "fastapi",
+        "isExtraImport": true,
+        "detail": "fastapi",
+        "documentation": {}
+    },
+    {
+        "label": "Form",
+        "importPath": "fastapi",
+        "description": "fastapi",
+        "isExtraImport": true,
+        "detail": "fastapi",
+        "documentation": {}
+    },
+    {
+        "label": "File",
+        "importPath": "fastapi",
+        "description": "fastapi",
+        "isExtraImport": true,
+        "detail": "fastapi",
+        "documentation": {}
+    },
+    {
+        "label": "JSONResponse",
+        "importPath": "fastapi.responses",
+        "description": "fastapi.responses",
+        "isExtraImport": true,
+        "detail": "fastapi.responses",
+        "documentation": {}
+    },
+    {
+        "label": "CORSMiddleware",
+        "importPath": "fastapi.middleware.cors",
+        "description": "fastapi.middleware.cors",
+        "isExtraImport": true,
+        "detail": "fastapi.middleware.cors",
+        "documentation": {}
+    },
+    {
+        "label": "llm",
+        "importPath": "src.config.llm",
+        "description": "src.config.llm",
+        "isExtraImport": true,
+        "detail": "src.config.llm",
+        "documentation": {}
+    },
+    {
+        "label": "format_prompt",
+        "importPath": "src.prompt.promt",
+        "description": "src.prompt.promt",
+        "isExtraImport": true,
+        "detail": "src.prompt.promt",
+        "documentation": {}
+    },
+    {
+        "label": "JsonOutputParser",
+        "importPath": "langchain_core.output_parsers",
+        "description": "langchain_core.output_parsers",
+        "isExtraImport": true,
+        "detail": "langchain_core.output_parsers",
+        "documentation": {}
+    },
+    {
+        "label": "uvicorn",
+        "kind": 6,
+        "isExtraImport": true,
+        "importPath": "uvicorn",
+        "description": "uvicorn",
+        "detail": "uvicorn",
+        "documentation": {}
+    },
+    {
+        "label": "BytesIO",
+        "importPath": "io",
+        "description": "io",
+        "isExtraImport": true,
+        "detail": "io",
+        "documentation": {}
+    },
+    {
+        "label": "base64",
+        "kind": 6,
+        "isExtraImport": true,
+        "importPath": "base64",
+        "description": "base64",
+        "detail": "base64",
+        "documentation": {}
+    },
+    {
+        "label": "Field",
+        "importPath": "pydantic",
+        "description": "pydantic",
+        "isExtraImport": true,
+        "detail": "pydantic",
+        "documentation": {}
+    },
+    {
+        "label": "BaseModel",
+        "importPath": "pydantic",
+        "description": "pydantic",
+        "isExtraImport": true,
+        "detail": "pydantic",
+        "documentation": {}
+    },
+    {
+        "label": "ThreadPoolExecutor",
+        "importPath": "concurrent.futures",
+        "description": "concurrent.futures",
+        "isExtraImport": true,
+        "detail": "concurrent.futures",
+        "documentation": {}
+    },
+    {
+        "label": "asyncio",
+        "kind": 6,
+        "isExtraImport": true,
+        "importPath": "asyncio",
+        "description": "asyncio",
+        "detail": "asyncio",
+        "documentation": {}
+    },
+    {
+        "label": "os",
+        "kind": 6,
+        "isExtraImport": true,
+        "importPath": "os",
+        "description": "os",
+        "detail": "os",
+        "documentation": {}
+    },
+    {
+        "label": "functools",
+        "kind": 6,
+        "isExtraImport": true,
+        "importPath": "functools",
+        "description": "functools",
+        "detail": "functools",
+        "documentation": {}
+    },
+    {
+        "label": "threading",
+        "kind": 6,
+        "isExtraImport": true,
+        "importPath": "threading",
+        "description": "threading",
+        "detail": "threading",
+        "documentation": {}
+    },
+    {
+        "label": "inference",
+        "importPath": "src.inference.segment_inference",
+        "description": "src.inference.segment_inference",
+        "isExtraImport": true,
+        "detail": "src.inference.segment_inference",
+        "documentation": {}
+    },
+    {
+        "label": "llm",
+        "kind": 5,
+        "importPath": "src.config.llm",
+        "description": "src.config.llm",
+        "peekOfCode": "llm = GoogleGenerativeAI(\n    model=\"gemini-1.5-flash\",\n    temperature=0,\n)",
+        "detail": "src.config.llm",
+        "documentation": {}
+    },
+    {
+        "label": "inference",
+        "kind": 2,
+        "importPath": "src.inference.segment_inference",
+        "description": "src.inference.segment_inference",
+        "peekOfCode": "def inference(image: np.array, threshold_confidence=0.5, threshold_iou=0.7):\n    input = preprocess(image)\n    outputs = postprocess(\n        model.run(None, {\"images\": input}),\n        threshold_confidence=threshold_confidence,\n        threshold_iou=threshold_iou,\n    )\n    return outputs",
+        "detail": "src.inference.segment_inference",
+        "documentation": {}
+    },
+    {
+        "label": "model_path",
+        "kind": 5,
+        "importPath": "src.inference.segment_inference",
+        "description": "src.inference.segment_inference",
+        "peekOfCode": "model_path = \"./src/model/segment.onnx\"\nmodel = ort.InferenceSession(\n    model_path,\n)\ndef inference(image: np.array, threshold_confidence=0.5, threshold_iou=0.7):\n    input = preprocess(image)\n    outputs = postprocess(\n        model.run(None, {\"images\": input}),\n        threshold_confidence=threshold_confidence,\n        threshold_iou=threshold_iou,",
+        "detail": "src.inference.segment_inference",
+        "documentation": {}
+    },
+    {
+        "label": "model",
+        "kind": 5,
+        "importPath": "src.inference.segment_inference",
+        "description": "src.inference.segment_inference",
+        "peekOfCode": "model = ort.InferenceSession(\n    model_path,\n)\ndef inference(image: np.array, threshold_confidence=0.5, threshold_iou=0.7):\n    input = preprocess(image)\n    outputs = postprocess(\n        model.run(None, {\"images\": input}),\n        threshold_confidence=threshold_confidence,\n        threshold_iou=threshold_iou,\n    )",
+        "detail": "src.inference.segment_inference",
+        "documentation": {}
+    },
+    {
+        "label": "format_prompt",
+        "kind": 5,
+        "importPath": "src.prompt.promt",
+        "description": "src.prompt.promt",
+        "peekOfCode": "format_prompt = \"\"\"\n#Role: You are an expert at correcting spelling errors from interviewee's resume information.\n#Instruction:\nYou are provided with a dictionary containing information from the user's resume by an OCR model. It may have misspellings or wrong entries.\nPlease correct the spelling of each field.\nMove the content of the fields to more appropriate fields if necessary.\nYou must not fabricate information and create new information.\nYou must return JSON containing the same format as the original format:\n#Input:\nMy resume is as follows: {input}",
+        "detail": "src.prompt.promt",
+        "documentation": {}
+    },
+    {
+        "label": "preprocess",
+        "kind": 2,
+        "importPath": "src.utils.utils_segment",
+        "description": "src.utils.utils_segment",
+        "peekOfCode": "def preprocess(img: np.array, shape=(640, 640)) -> np.array:\n    global img_width, img_height, left, top, ratio\n    img, ratio, (left, top) = resize_and_pad(img, new_shape=shape)\n    img_height, img_width, _ = img.shape\n    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\n    img = img.transpose(2, 0, 1)\n    img = img.reshape(1, 3, 640, 640).astype(\"float32\")\n    img = img / 255.0\n    return img\ndef extract_box(outputs):",
+        "detail": "src.utils.utils_segment",
+        "documentation": {}
+    },
+    {
+        "label": "extract_box",
+        "kind": 2,
+        "importPath": "src.utils.utils_segment",
+        "description": "src.utils.utils_segment",
+        "peekOfCode": "def extract_box(outputs):\n    output0 = outputs[0]\n    output1 = outputs[1]\n    output0 = output0[0].transpose()\n    output1 = output1[0]\n    boxes = output0[:, 0:number_class_custom]\n    masks = output0[:, number_class_custom:]\n    output1 = output1.reshape(32, 160 * 160)\n    output1 = output1.reshape(32, 160 * 160)\n    masks = masks @ output1",
+        "detail": "src.utils.utils_segment",
+        "documentation": {}
+    },
+    {
+        "label": "intersection",
+        "kind": 2,
+        "importPath": "src.utils.utils_segment",
+        "description": "src.utils.utils_segment",
+        "peekOfCode": "def intersection(box1, box2):\n    box1_x1, box1_y1, box1_x2, box1_y2 = box1[:4]\n    box2_x1, box2_y1, box2_x2, box2_y2 = box2[:4]\n    x1 = max(box1_x1, box2_x1)\n    y1 = max(box1_y1, box2_y1)\n    x2 = min(box1_x2, box2_x2)\n    y2 = min(box1_y2, box2_y2)\n    return (x2 - x1) * (y2 - y1)\ndef union(box1, box2):\n    box1_x1, box1_y1, box1_x2, box1_y2 = box1[:4]",
+        "detail": "src.utils.utils_segment",
+        "documentation": {}
+    },
+    {
+        "label": "union",
+        "kind": 2,
+        "importPath": "src.utils.utils_segment",
+        "description": "src.utils.utils_segment",
+        "peekOfCode": "def union(box1, box2):\n    box1_x1, box1_y1, box1_x2, box1_y2 = box1[:4]\n    box2_x1, box2_y1, box2_x2, box2_y2 = box2[:4]\n    box1_area = (box1_x2 - box1_x1) * (box1_y2 - box1_y1)\n    box2_area = (box2_x2 - box2_x1) * (box2_y2 - box2_y1)\n    return box1_area + box2_area - intersection(box1, box2)\ndef iou(box1, box2):\n    return intersection(box1, box2) / union(box1, box2)\ndef sigmoid(z):\n    return 1 / (1 + np.exp(-z))",
+        "detail": "src.utils.utils_segment",
+        "documentation": {}
+    },
+    {
+        "label": "iou",
+        "kind": 2,
+        "importPath": "src.utils.utils_segment",
+        "description": "src.utils.utils_segment",
+        "peekOfCode": "def iou(box1, box2):\n    return intersection(box1, box2) / union(box1, box2)\ndef sigmoid(z):\n    return 1 / (1 + np.exp(-z))\ndef get_mask(row, box, img_width, img_height, threshold):\n    mask = row.reshape(160, 160)\n    mask = sigmoid(mask)\n    mask = (mask > threshold).astype(\"uint8\") * 255\n    x1, y1, x2, y2 = box\n    mask_x1 = round(x1 / img_width * 160)",
+        "detail": "src.utils.utils_segment",
+        "documentation": {}
+    },
+    {
+        "label": "sigmoid",
+        "kind": 2,
+        "importPath": "src.utils.utils_segment",
+        "description": "src.utils.utils_segment",
+        "peekOfCode": "def sigmoid(z):\n    return 1 / (1 + np.exp(-z))\ndef get_mask(row, box, img_width, img_height, threshold):\n    mask = row.reshape(160, 160)\n    mask = sigmoid(mask)\n    mask = (mask > threshold).astype(\"uint8\") * 255\n    x1, y1, x2, y2 = box\n    mask_x1 = round(x1 / img_width * 160)\n    mask_y1 = round(y1 / img_height * 160)\n    mask_x2 = round(x2 / img_width * 160)",
+        "detail": "src.utils.utils_segment",
+        "documentation": {}
+    },
+    {
+        "label": "get_mask",
+        "kind": 2,
+        "importPath": "src.utils.utils_segment",
+        "description": "src.utils.utils_segment",
+        "peekOfCode": "def get_mask(row, box, img_width, img_height, threshold):\n    mask = row.reshape(160, 160)\n    mask = sigmoid(mask)\n    mask = (mask > threshold).astype(\"uint8\") * 255\n    x1, y1, x2, y2 = box\n    mask_x1 = round(x1 / img_width * 160)\n    mask_y1 = round(y1 / img_height * 160)\n    mask_x2 = round(x2 / img_width * 160)\n    mask_y2 = round(y2 / img_height * 160)\n    mask = mask[mask_y1:mask_y2, mask_x1:mask_x2]",
+        "detail": "src.utils.utils_segment",
+        "documentation": {}
+    },
+    {
+        "label": "get_polygon",
+        "kind": 2,
+        "importPath": "src.utils.utils_segment",
+        "description": "src.utils.utils_segment",
+        "peekOfCode": "def get_polygon(mask):\n    contours = cv2.findContours(mask, cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE)\n    polygon = [[contour[0][0], contour[0][1]] for contour in contours[0][0]]\n    return polygon\ndef postprocess(outputs, threshold_confidence, threshold_iou):\n    objects = []\n    for row in extract_box(outputs):\n        xc, yc, w, h = row[:4]\n        x1 = (xc - w / 2) / 640 * img_width\n        y1 = (yc - h / 2) / 640 * img_height",
+        "detail": "src.utils.utils_segment",
+        "documentation": {}
+    },
+    {
+        "label": "postprocess",
+        "kind": 2,
+        "importPath": "src.utils.utils_segment",
+        "description": "src.utils.utils_segment",
+        "peekOfCode": "def postprocess(outputs, threshold_confidence, threshold_iou):\n    objects = []\n    for row in extract_box(outputs):\n        xc, yc, w, h = row[:4]\n        x1 = (xc - w / 2) / 640 * img_width\n        y1 = (yc - h / 2) / 640 * img_height\n        x2 = (xc + w / 2) / 640 * img_width\n        y2 = (yc + h / 2) / 640 * img_height\n        prob = row[4:number_class_custom].max()\n        if prob < threshold_confidence:",
+        "detail": "src.utils.utils_segment",
+        "documentation": {}
+    },
+    {
+        "label": "extract_text_dict",
+        "kind": 2,
+        "importPath": "src.utils.utils_segment",
+        "description": "src.utils.utils_segment",
+        "peekOfCode": "def extract_text_dict(outputs):\n    result_dict = {}\n    for output in outputs:\n        label = output.get(\"label\").lower()\n        text = output.get(\"text\")\n        if label in result_dict:\n            result_dict[label] += \" \" + text\n        else:\n            result_dict[label] = text\n    return result_dict",
+        "detail": "src.utils.utils_segment",
+        "documentation": {}
+    },
+    {
+        "label": "extract_text",
+        "kind": 2,
+        "importPath": "src.utils.utils_segment",
+        "description": "src.utils.utils_segment",
+        "peekOfCode": "def extract_text(outputs, image_origin):\n    for i in range(len(outputs)):\n        image = crop_image(image_origin, outputs[i].get(\"box\"))\n        text = pytesseract.image_to_string(image)\n        outputs[i].update({\"text\": text})\n        if \"text\" in outputs[i]:\n            outputs[i][\"text\"] += text\n        else:\n            outputs[i].update({\"text\": text})\n    return extract_text_dict(outputs)",
+        "detail": "src.utils.utils_segment",
+        "documentation": {}
+    },
+    {
+        "label": "crop_image",
+        "kind": 2,
+        "importPath": "src.utils.utils_segment",
+        "description": "src.utils.utils_segment",
+        "peekOfCode": "def crop_image(image, box):\n    x1, y1, x2, y2 = map(int, box)\n    cropped_image = image[y1:y2, x1:x2]\n    return cropped_image\ndef resize_and_pad(\n    image: np.array,\n    new_shape: Tuple[int, int],\n    padding_color: Tuple[int] = (144, 144, 144),\n) -> np.array:\n    h_org, w_org = image.shape[:2]",
+        "detail": "src.utils.utils_segment",
+        "documentation": {}
+    },
+    {
+        "label": "resize_and_pad",
+        "kind": 2,
+        "importPath": "src.utils.utils_segment",
+        "description": "src.utils.utils_segment",
+        "peekOfCode": "def resize_and_pad(\n    image: np.array,\n    new_shape: Tuple[int, int],\n    padding_color: Tuple[int] = (144, 144, 144),\n) -> np.array:\n    h_org, w_org = image.shape[:2]\n    w_new, h_new = new_shape\n    padd_left, padd_right, padd_top, padd_bottom = 0, 0, 0, 0\n    # Padding left to right\n    if h_org >= w_org:",
+        "detail": "src.utils.utils_segment",
+        "documentation": {}
+    },
+    {
+        "label": "unpad_and_resize_boxes",
+        "kind": 2,
+        "importPath": "src.utils.utils_segment",
+        "description": "src.utils.utils_segment",
+        "peekOfCode": "def unpad_and_resize_boxes(boxes, ratio, left, top):\n    if len(boxes) == 0:\n        return boxes\n    boxes = np.array(boxes)\n    if boxes.ndim == 1:\n        boxes = boxes.reshape(-1, 4)\n    boxes[:, [0, 2]] -= left\n    boxes[:, [1, 3]] -= top\n    boxes[:, :4] /= ratio\n    if len(boxes) == 1:",
+        "detail": "src.utils.utils_segment",
+        "documentation": {}
+    },
+    {
+        "label": "draw_bounding_boxes",
+        "kind": 2,
+        "importPath": "src.utils.utils_segment",
+        "description": "src.utils.utils_segment",
+        "peekOfCode": "def draw_bounding_boxes(image, outputs):\n    # Create a copy of the image to draw on\n    image_with_boxes = image.copy()\n    # Define a list of colors for the bounding boxes\n    label_colors = {\n        \"Certifications\": (255, 0, 0),\n        \"Community\": (0, 255, 0),\n        \"Contact\": (0, 0, 255),\n        \"Education\": (255, 128, 0),\n        \"Experience\": (255, 0, 255),",
+        "detail": "src.utils.utils_segment",
+        "documentation": {}
+    },
+    {
+        "label": "class_names",
+        "kind": 5,
+        "importPath": "src.utils.utils_segment",
+        "description": "src.utils.utils_segment",
+        "peekOfCode": "class_names = [\n    \"Certifications\",\n    \"Community\",\n    \"Contact\",\n    \"Education\",\n    \"Experience\",\n    \"Interests\",\n    \"Languages\",\n    \"Name\",\n    \"Profile\",",
+        "detail": "src.utils.utils_segment",
+        "documentation": {}
+    },
+    {
+        "label": "number_class_custom",
+        "kind": 5,
+        "importPath": "src.utils.utils_segment",
+        "description": "src.utils.utils_segment",
+        "peekOfCode": "number_class_custom = int(len(class_names) + 4)\nimg_width, img_height = None, None\nleft = None\ntop = None\nratio = None\ndef preprocess(img: np.array, shape=(640, 640)) -> np.array:\n    global img_width, img_height, left, top, ratio\n    img, ratio, (left, top) = resize_and_pad(img, new_shape=shape)\n    img_height, img_width, _ = img.shape\n    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)",
+        "detail": "src.utils.utils_segment",
+        "documentation": {}
+    },
+    {
+        "label": "left",
+        "kind": 5,
+        "importPath": "src.utils.utils_segment",
+        "description": "src.utils.utils_segment",
+        "peekOfCode": "left = None\ntop = None\nratio = None\ndef preprocess(img: np.array, shape=(640, 640)) -> np.array:\n    global img_width, img_height, left, top, ratio\n    img, ratio, (left, top) = resize_and_pad(img, new_shape=shape)\n    img_height, img_width, _ = img.shape\n    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\n    img = img.transpose(2, 0, 1)\n    img = img.reshape(1, 3, 640, 640).astype(\"float32\")",
+        "detail": "src.utils.utils_segment",
+        "documentation": {}
+    },
+    {
+        "label": "top",
+        "kind": 5,
+        "importPath": "src.utils.utils_segment",
+        "description": "src.utils.utils_segment",
+        "peekOfCode": "top = None\nratio = None\ndef preprocess(img: np.array, shape=(640, 640)) -> np.array:\n    global img_width, img_height, left, top, ratio\n    img, ratio, (left, top) = resize_and_pad(img, new_shape=shape)\n    img_height, img_width, _ = img.shape\n    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\n    img = img.transpose(2, 0, 1)\n    img = img.reshape(1, 3, 640, 640).astype(\"float32\")\n    img = img / 255.0",
+        "detail": "src.utils.utils_segment",
+        "documentation": {}
+    },
+    {
+        "label": "ratio",
+        "kind": 5,
+        "importPath": "src.utils.utils_segment",
+        "description": "src.utils.utils_segment",
+        "peekOfCode": "ratio = None\ndef preprocess(img: np.array, shape=(640, 640)) -> np.array:\n    global img_width, img_height, left, top, ratio\n    img, ratio, (left, top) = resize_and_pad(img, new_shape=shape)\n    img_height, img_width, _ = img.shape\n    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\n    img = img.transpose(2, 0, 1)\n    img = img.reshape(1, 3, 640, 640).astype(\"float32\")\n    img = img / 255.0\n    return img",
+        "detail": "src.utils.utils_segment",
+        "documentation": {}
+    },
+    {
+        "label": "LLMRequest",
+        "kind": 6,
+        "importPath": "app",
+        "description": "app",
+        "peekOfCode": "class LLMRequest(BaseModel):\n    text: str = Field(..., title=\"Text to generate completion\")\ndef call_llm(data):\n    input = format_prompt.format(input=data)\n    response = llm.invoke(input)\n    response = JsonOutputParser().parse(response)\n    return response\n@app.post(\"/llm\", status_code=status.HTTP_200_OK)\nasync def llm_predict(data: LLMRequest):\n    try:",
+        "detail": "app",
+        "documentation": {}
+    },
+    {
+        "label": "run_in_thread",
+        "kind": 2,
+        "importPath": "app",
+        "description": "app",
+        "peekOfCode": "def run_in_thread(func, *args, **kwargs):\n    loop = asyncio.get_event_loop()\n    func_name = func.__name__\n    def wrapper(*args, **kwargs):\n        thread_id = threading.get_ident()\n        print(f\"[Running function '{func_name}' in thread ID: {thread_id}]\")\n        return func(*args, **kwargs)\n    return loop.run_in_executor(executor, functools.partial(wrapper, *args, **kwargs))\ndef predict_func(threshold_confidence, threshold_iou, image):\n    image = np.frombuffer(image, np.uint8)",
+        "detail": "app",
+        "documentation": {}
+    },
+    {
+        "label": "predict_func",
+        "kind": 2,
+        "importPath": "app",
+        "description": "app",
+        "peekOfCode": "def predict_func(threshold_confidence, threshold_iou, image):\n    image = np.frombuffer(image, np.uint8)\n    image = cv2.imdecode(image, cv2.IMREAD_COLOR)\n    outputs = inference(\n        image,\n        threshold_confidence=threshold_confidence,\n        threshold_iou=threshold_iou,\n    )\n    text = extract_text(outputs=outputs, image_origin=image)\n    image = draw_bounding_boxes(image, outputs)",
+        "detail": "app",
+        "documentation": {}
+    },
+    {
+        "label": "call_llm",
+        "kind": 2,
+        "importPath": "app",
+        "description": "app",
+        "peekOfCode": "def call_llm(data):\n    input = format_prompt.format(input=data)\n    response = llm.invoke(input)\n    response = JsonOutputParser().parse(response)\n    return response\n@app.post(\"/llm\", status_code=status.HTTP_200_OK)\nasync def llm_predict(data: LLMRequest):\n    try:\n        response = await run_in_thread(call_llm, data.text)\n        return JSONResponse(content=response, status_code=status.HTTP_200_OK)",
+        "detail": "app",
+        "documentation": {}
+    },
+    {
+        "label": "app",
+        "kind": 5,
+        "importPath": "app",
+        "description": "app",
+        "peekOfCode": "app = FastAPI(docs_url=\"/\")\napp.add_middleware(\n    CORSMiddleware,\n    allow_origins=[\"*\"],\n    allow_credentials=True,\n    allow_methods=[\"*\"],\n    allow_headers=[\"*\"],\n)\nexecutor = ThreadPoolExecutor(max_workers=int(os.cpu_count() + 4))\ndef run_in_thread(func, *args, **kwargs):",
+        "detail": "app",
+        "documentation": {}
+    },
+    {
+        "label": "executor",
+        "kind": 5,
+        "importPath": "app",
+        "description": "app",
+        "peekOfCode": "executor = ThreadPoolExecutor(max_workers=int(os.cpu_count() + 4))\ndef run_in_thread(func, *args, **kwargs):\n    loop = asyncio.get_event_loop()\n    func_name = func.__name__\n    def wrapper(*args, **kwargs):\n        thread_id = threading.get_ident()\n        print(f\"[Running function '{func_name}' in thread ID: {thread_id}]\")\n        return func(*args, **kwargs)\n    return loop.run_in_executor(executor, functools.partial(wrapper, *args, **kwargs))\ndef predict_func(threshold_confidence, threshold_iou, image):",
+        "detail": "app",
+        "documentation": {}
+    }
+]

requirements.txt

@@ -6,4 +6,6 @@ langchain-google-genai
 uvicorn
 python-multipart
 asyncio
-onnxruntime
+onnxruntime
+langgraph
+langchain_community

src/config/llm.py

@@ -1,5 +1,6 @@
-from langchain_google_genai import GoogleGenerativeAI
-llm = GoogleGenerativeAI(
+from langchain_google_genai import ChatGoogleGenerativeAI
+
+llm = ChatGoogleGenerativeAI(
     model="gemini-1.5-flash",
     temperature=0,
 )

src/graph.py

@@ -0,0 +1,117 @@
+from langgraph.graph import StateGraph, END, START, add_messages
+from typing import TypedDict, Any, Annotated
+from PIL import Image
+from src.utils.utils_segment import (
+    preprocess,
+    postprocess,
+    extract_text,
+    draw_bounding_boxes,
+)
+from src.inference.segment_inference import model
+from src.config.llm import llm
+from src.prompt.promt import format_prompt
+from langchain_core.output_parsers import JsonOutputParser
+
+parser = JsonOutputParser()
+
+
+class State(TypedDict):
+    image: Any
+    image_origin: Any
+    outputs_from_inference: Any
+    text_extracted_from_ocr: Any
+    threshold_confidence: float
+    threshold_iou: float
+    cropped_images: Any
+    parser_output: bool
+    image_with_bounding_boxes: Any
+    _image: Annotated[Any, add_messages]
+    crop_image: Any
+
+
+class N:
+    PRE_PROCESS = "PRE_PROCESS"
+    POST_PROCESS = "POST_PROCESS"
+    INFERENCE = "INFERENCE"
+    EXTRACT_TEXT_FROM_OCR = "EXTRACT_TEXT_FROM_OCR"
+    PARSER_WITH_LLM = "PARSER_WITH_LLM"
+    IMAGE_WITH_BOUNDING_BOXES = "IMAGE_WITH_BOUNDING_BOXES"
+
+
+workflow = StateGraph(State)
+
+
+def pre_process_fn(state: State):
+    preprocess_img = preprocess(state["image_origin"])
+    print("preprocess_img", preprocess_img.shape)
+    image_for_display = (preprocess_img[0] * 255).astype("uint8")
+    image_for_display = image_for_display.transpose(1, 2, 0)
+    image_show = Image.fromarray(image_for_display)
+    return {"image": preprocess_img, "_image": image_show}
+
+
+def inference_fn(state: State):
+    image = state["image"]
+    outputs = model.run(None, {"images": image})
+    return {"outputs_from_inference": outputs}
+
+
+def post_process_fn(state: State):
+    outputs = state["outputs_from_inference"]
+    threshold_confidence = state["threshold_confidence"]
+    threshold_iou = state["threshold_iou"]
+    post_process_output = postprocess(outputs, threshold_confidence, threshold_iou)
+
+    return {
+        "outputs_from_inference": post_process_output,
+    }
+
+
+def extract_text_from_ocr_fn(state: State):
+    image_origin = state["image_origin"]
+    output_from_inference = state["outputs_from_inference"]
+
+    text = extract_text(output_from_inference, image_origin)
+    return {"text_extracted_from_ocr": text}
+
+
+def draw_bounding_boxes_fn(state: State):
+    image = state["image_origin"]
+    outputs = state["outputs_from_inference"]
+    image_with_bounding_boxes = draw_bounding_boxes(image, outputs)
+    return {"image_with_bounding_boxes": image_with_bounding_boxes}
+
+
+def parser_output_fn(state: State):
+    text_extracted_from_ocr = state["text_extracted_from_ocr"]
+    chain = format_prompt | llm | parser
+    response = chain.invoke({"user_input": text_extracted_from_ocr})
+    print(response)
+    return {"parser_output": response}
+
+#NODE
+workflow.add_node(N.PRE_PROCESS, pre_process_fn)
+workflow.add_node(N.INFERENCE, inference_fn)
+workflow.add_node(N.POST_PROCESS, post_process_fn)
+workflow.add_node(N.EXTRACT_TEXT_FROM_OCR, extract_text_from_ocr_fn)
+workflow.add_node(N.IMAGE_WITH_BOUNDING_BOXES, draw_bounding_boxes_fn)
+workflow.add_node(N.PARSER_WITH_LLM, parser_output_fn)
+
+#EDGE
+workflow.add_edge(START, N.PRE_PROCESS)
+workflow.add_edge(N.PRE_PROCESS, N.INFERENCE)
+workflow.add_edge(N.INFERENCE, N.POST_PROCESS)
+workflow.add_edge(N.POST_PROCESS, N.IMAGE_WITH_BOUNDING_BOXES)
+workflow.add_edge(N.IMAGE_WITH_BOUNDING_BOXES, N.EXTRACT_TEXT_FROM_OCR)
+workflow.add_conditional_edges(
+    N.EXTRACT_TEXT_FROM_OCR,
+    lambda state: N.PARSER_WITH_LLM if state["parser_output"] else END,
+    {
+        N.PARSER_WITH_LLM: N.PARSER_WITH_LLM,
+        END: END,
+    },
+)
+workflow.add_edge(N.PARSER_WITH_LLM, END)
+
+
+app = workflow.compile()

src/inference/segment_inference.py

@@ -2,7 +2,7 @@ import onnxruntime as ort
 from src.utils.utils_segment import preprocess, postprocess
 import numpy as np
 
-model_path = "./src/model/segment.onnx"
+model_path = "./src/model/best.onnx"
 model = ort.InferenceSession(
     model_path,
 )

src/model/{segment.onnx → best.onnx}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:e38660d6cb501bc21d33249f1e1dffd9038d57a82f27bd089746e1aa8eca53a9
-size 109150130
+oid sha256:d1fe98385269d978abe302212d7eb38d0acd584eaa3af6e45ddfbe082a307f7e
+size 109140800

src/notebook/notebook.ipynb

@@ -2,22 +2,65 @@
  "cells": [
   {
    "cell_type": "code",
-   "execution_count": 12,
+   "execution_count": 1,
    "metadata": {},
    "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Requirement already satisfied: roboflow in c:\\users\\htbqn\\appdata\\local\\programs\\python\\python311\\lib\\site-packages (1.1.44)\n",
+      "Requirement already satisfied: certifi in c:\\users\\htbqn\\appdata\\local\\programs\\python\\python311\\lib\\site-packages (from roboflow) (2024.7.4)\n",
+      "Requirement already satisfied: idna==3.7 in c:\\users\\htbqn\\appdata\\local\\programs\\python\\python311\\lib\\site-packages (from roboflow) (3.7)\n",
+      "Requirement already satisfied: cycler in c:\\users\\htbqn\\appdata\\local\\programs\\python\\python311\\lib\\site-packages (from roboflow) (0.12.1)\n",
+      "Requirement already satisfied: kiwisolver>=1.3.1 in c:\\users\\htbqn\\appdata\\local\\programs\\python\\python311\\lib\\site-packages (from roboflow) (1.4.5)\n",
+      "Requirement already satisfied: matplotlib in c:\\users\\htbqn\\appdata\\local\\programs\\python\\python311\\lib\\site-packages (from roboflow) (3.9.2)\n",
+      "Requirement already satisfied: numpy>=1.18.5 in c:\\users\\htbqn\\appdata\\local\\programs\\python\\python311\\lib\\site-packages (from roboflow) (1.26.4)\n",
+      "Requirement already satisfied: opencv-python-headless==4.10.0.84 in c:\\users\\htbqn\\appdata\\local\\programs\\python\\python311\\lib\\site-packages (from roboflow) (4.10.0.84)\n",
+      "Requirement already satisfied: Pillow>=7.1.2 in c:\\users\\htbqn\\appdata\\local\\programs\\python\\python311\\lib\\site-packages (from roboflow) (10.4.0)\n",
+      "Requirement already satisfied: python-dateutil in c:\\users\\htbqn\\appdata\\local\\programs\\python\\python311\\lib\\site-packages (from roboflow) (2.9.0.post0)\n",
+      "Requirement already satisfied: python-dotenv in c:\\users\\htbqn\\appdata\\local\\programs\\python\\python311\\lib\\site-packages (from roboflow) (1.0.1)\n",
+      "Requirement already satisfied: requests in c:\\users\\htbqn\\appdata\\local\\programs\\python\\python311\\lib\\site-packages (from roboflow) (2.31.0)\n",
+      "Requirement already satisfied: six in c:\\users\\htbqn\\appdata\\local\\programs\\python\\python311\\lib\\site-packages (from roboflow) (1.16.0)\n",
+      "Requirement already satisfied: urllib3>=1.26.6 in c:\\users\\htbqn\\appdata\\local\\programs\\python\\python311\\lib\\site-packages (from roboflow) (2.2.3)\n",
+      "Requirement already satisfied: tqdm>=4.41.0 in c:\\users\\htbqn\\appdata\\local\\programs\\python\\python311\\lib\\site-packages (from roboflow) (4.66.4)\n",
+      "Requirement already satisfied: PyYAML>=5.3.1 in c:\\users\\htbqn\\appdata\\local\\programs\\python\\python311\\lib\\site-packages (from roboflow) (6.0.1)\n",
+      "Requirement already satisfied: requests-toolbelt in c:\\users\\htbqn\\appdata\\local\\programs\\python\\python311\\lib\\site-packages (from roboflow) (1.0.0)\n",
+      "Requirement already satisfied: filetype in c:\\users\\htbqn\\appdata\\local\\programs\\python\\python311\\lib\\site-packages (from roboflow) (1.2.0)\n",
+      "Requirement already satisfied: colorama in c:\\users\\htbqn\\appdata\\local\\programs\\python\\python311\\lib\\site-packages (from tqdm>=4.41.0->roboflow) (0.4.4)\n",
+      "Requirement already satisfied: contourpy>=1.0.1 in c:\\users\\htbqn\\appdata\\local\\programs\\python\\python311\\lib\\site-packages (from matplotlib->roboflow) (1.2.1)\n",
+      "Requirement already satisfied: fonttools>=4.22.0 in c:\\users\\htbqn\\appdata\\local\\programs\\python\\python311\\lib\\site-packages (from matplotlib->roboflow) (4.53.1)\n",
+      "Requirement already satisfied: packaging>=20.0 in c:\\users\\htbqn\\appdata\\local\\programs\\python\\python311\\lib\\site-packages (from matplotlib->roboflow) (24.1)\n",
+      "Requirement already satisfied: pyparsing>=2.3.1 in c:\\users\\htbqn\\appdata\\local\\programs\\python\\python311\\lib\\site-packages (from matplotlib->roboflow) (3.1.2)\n",
+      "Requirement already satisfied: charset-normalizer<4,>=2 in c:\\users\\htbqn\\appdata\\local\\programs\\python\\python311\\lib\\site-packages (from requests->roboflow) (3.3.2)\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "WARNING: Ignoring invalid distribution ~angchain (C:\\Users\\htbqn\\AppData\\Local\\Programs\\Python\\Python311\\Lib\\site-packages)\n",
+      "WARNING: Ignoring invalid distribution ~ip (C:\\Users\\htbqn\\AppData\\Local\\Programs\\Python\\Python311\\Lib\\site-packages)\n",
+      "WARNING: Ignoring invalid distribution ~angchain (C:\\Users\\htbqn\\AppData\\Local\\Programs\\Python\\Python311\\Lib\\site-packages)\n",
+      "WARNING: Ignoring invalid distribution ~ip (C:\\Users\\htbqn\\AppData\\Local\\Programs\\Python\\Python311\\Lib\\site-packages)\n",
+      "WARNING: Ignoring invalid distribution ~angchain (C:\\Users\\htbqn\\AppData\\Local\\Programs\\Python\\Python311\\Lib\\site-packages)\n",
+      "WARNING: Ignoring invalid distribution ~ip (C:\\Users\\htbqn\\AppData\\Local\\Programs\\Python\\Python311\\Lib\\site-packages)\n"
+     ]
+    },
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
       "loading Roboflow workspace...\n",
-      "loading Roboflow project...\n"
+      "loading Roboflow project...\n",
+      "Dependency ultralytics==8.0.196 is required but found version=8.2.90, to fix: `pip install ultralytics==8.0.196`\n"
      ]
     },
     {
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "Downloading Dataset Version Zip in cvparsing-2 to yolov9:: 100%|██████████| 63864/63864 [00:04<00:00, 15236.33it/s]"
+      "Downloading Dataset Version Zip in annotation-2 to yolov8:: 100%|██████████| 92659/92659 [00:06<00:00, 14429.98it/s]"
      ]
     },
     {
@@ -32,7 +75,7 @@
      "output_type": "stream",
      "text": [
       "\n",
-      "Extracting Dataset Version Zip to cvparsing-2 in yolov9:: 100%|██████████| 2344/2344 [00:00<00:00, 5118.00it/s]\n"
+      "Extracting Dataset Version Zip to annotation-2 in yolov8:: 100%|██████████| 3376/3376 [00:00<00:00, 3563.34it/s]\n"
      ]
     }
    ],
@@ -104,7 +147,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 6,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -115,6 +158,10 @@
     "nc: 14\n",
     "names: ['Achievement', 'Certifications', 'Community', 'Contact', 'Education', 'Experience', 'Interests', 'Languages', 'Name', 'Profil', 'Projects', 'image', 'resume', 'skills']\"\"\"\n",
     "\n",
+    "#  3 + 8  -> 3\n",
+    "#  5 + 10 + 1 -> 5\n",
+    "# 13 + 7 -> 13\n",
+    "\n",
     "with open(\"./data.yaml\", 'w') as file:\n",
     "    file.write(yaml_text),\n",
     "\n",
@@ -122,6 +169,77 @@
     "# %cat /kaggle/working/data.yaml\n"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Labels mapped and saved to ./test.txt\n"
+     ]
+    }
+   ],
+   "source": [
+    "import os\n",
+    "\n",
+    "def map_yolo_labels(file_path, label_mapping, output_file=None):\n",
+    "    \"\"\"\n",
+    "    Remaps the class labels in a YOLOv8 label file and saves the updated labels.\n",
+    "    \n",
+    "    Parameters:\n",
+    "    - file_path: Path to the YOLOv8 label file.\n",
+    "    - label_mapping: Dictionary for remapping labels (e.g., {3: 3, 8: 3, 5: 5, 10: 5, 1: 5, 13: 13, 7: 13}).\n",
+    "    - output_file: Path to save the remapped labels (optional). If not provided, it overwrites the original file.\n",
+    "    \n",
+    "    Returns:\n",
+    "    - None\n",
+    "    \"\"\"\n",
+    "    with open(file_path, 'r') as file:\n",
+    "        lines = file.readlines()\n",
+    "\n",
+    "    # Process each line and remap the class label\n",
+    "    updated_lines = []\n",
+    "    for line in lines:\n",
+    "        parts = line.strip().split()\n",
+    "        class_id = int(parts[0])  # The first part is the class label\n",
+    "        \n",
+    "        # Remap the class label using the mapping\n",
+    "        new_class_id = label_mapping.get(class_id, class_id)\n",
+    "        \n",
+    "        # Reconstruct the line with the new class label\n",
+    "        updated_line = f\"{new_class_id} {' '.join(parts[1:])}\\n\"\n",
+    "        updated_lines.append(updated_line)\n",
+    "    \n",
+    "    # Write to the output file (overwrite original file if output_file is not provided)\n",
+    "    if output_file is None:\n",
+    "        output_file = file_path\n",
+    "\n",
+    "    with open(output_file, 'w') as file:\n",
+    "        file.writelines(updated_lines)\n",
+    "    \n",
+    "    print(f\"Labels mapped and saved to {output_file}\")\n",
+    "\n",
+    "# Example usage:\n",
+    "label_mapping = {\n",
+    "    3: 3,    # Map 3 -> 3\n",
+    "    8: 3,    # Map 8 -> 3\n",
+    "    5: 5,    # Map 5 -> 5\n",
+    "    10: 5,   # Map 10 -> 5\n",
+    "    1: 5,    # Map 1 -> 5\n",
+    "    13: 13,  # Map 13 -> 13\n",
+    "    7: 13    # Map 7 -> 13\n",
+    "}\n",
+    "\n",
+    "# Path to your YOLOv8 label file\n",
+    "label_file = \"./test.txt\"\n",
+    "\n",
+    "# Call the function to remap labels\n",
+    "map_yolo_labels(label_file, label_mapping)\n"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": 20,
@@ -433,7 +551,7 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": ".venv",
+   "display_name": "Python 3",
    "language": "python",
    "name": "python3"
   },

src/prompt/promt.py

@@ -1,13 +1,56 @@
-format_prompt = """
-#Role: You are an expert at correcting spelling errors from interviewee's resume information.
-#Instruction:
+from langchain_core.prompts import ChatPromptTemplate
+
+format_prompt = ChatPromptTemplate.from_messages(
+    [
+        (
+            "system",
+            """
+Role: You are an expert at extracting key information about projects from the "experience" section of an OCR'd resume.
+
+Instruction:
 You are provided with a dictionary containing information from the user's resume by an OCR model. It may have misspellings or wrong entries.
 Please correct the spelling of each field.
 Move the content of the fields to more appropriate fields if necessary.
 You must not fabricate information and create new information.
 
-You must return JSON containing the same format as the original format:
+Output must be in JSON format following the same structure as the input.
+""",
+        ),
+        ("human", "{user_input}"),
+    ]
+)
+
+
+prompt_experience = ChatPromptTemplate.from_messages(
+    [
+        (
+            "system",
+            """
+# Role: You are an expert at extracting key information about projects from the "experience" section of an OCR'd resume.
+
+# Instruction:
+You are given a JSON object containing extracted resume data from an OCR model. This data likely contains errors like misspellings, merged words, and extracted noise. Your task is to:
+
+1. **Pre-process the "experience" field:**
+    * Correct misspellings using your knowledge of common resume terms and English vocabulary.
+    * Separate merged words and remove any obvious OCR noise.
+
+2. **Identify the "experience" field:** Locate the field labeled "experience" (or a similar label) within the JSON object.
+
+3. **Extract project information:** 
+    * Identify project mentions: Look for keywords and phrases that indicate a project, such as "project," "developed," "implemented," "designed," "contributed to," etc.
+    * Extract project details: For each project mentioned:
+        * Project name or description
+        * Role and contributions
+        * Technologies used
+        * Outcomes and achievements
+
+4. **Structure the output:** Return a JSON object with a "projects" field containing an array of extracted project details. The exact format can be flexible to accommodate variations in the input data.
+
+
 
-#Input:
-My resume is as follows: {input}
-"""
+""",
+        ),
+        ("human", "{user_input}"),
+    ]
+)

src/utils/utils_segment.py

@@ -4,19 +4,15 @@ import cv2
 from typing import Tuple
 from pytesseract import pytesseract
 
-# path_to_tesseract = r"C:\Program Files\Tesseract-OCR\tesseract.exe"
-# pytesseract.tesseract_cmd = path_to_tesseract
+path_to_tesseract = r"C:\Program Files\Tesseract-OCR\tesseract.exe"
+pytesseract.tesseract_cmd = path_to_tesseract
 class_names = [
-    "Certifications",
     "Community",
     "Contact",
     "Education",
     "Experience",
     "Interests",
-    "Languages",
-    "Name",
     "Profile",
-    "Projects",
     "Skills",
 ]
 number_class_custom = int(len(class_names) + 4)
@@ -232,22 +228,18 @@ def unpad_and_resize_boxes(boxes, ratio, left, top):
         return boxes.tolist()
 
 
-def draw_bounding_boxes(image, outputs):
+def draw_bounding_boxes(image, outputs, save_path="output_image.jpg"):
     # Create a copy of the image to draw on
     image_with_boxes = image.copy()
 
     # Define a list of colors for the bounding boxes
     label_colors = {
-        "Certifications": (255, 0, 0),
         "Community": (0, 255, 0),
         "Contact": (0, 0, 255),
         "Education": (255, 128, 0),
         "Experience": (255, 0, 255),
         "Interests": (128, 128, 128),
-        "Languages": (128, 0, 0),
-        "Name": (0, 128, 0),
         "Profile": (0, 0, 128),
-        "Projects": (128, 128, 0),
         "Skills": (128, 0, 128),
     }
 
@@ -255,7 +247,6 @@ def draw_bounding_boxes(image, outputs):
     for output in outputs:
         box = output["box"]
         label = output["label"]
-        text = output.get("text", "")
 
         # Get the color for the label
         color = label_colors.get(
@@ -279,5 +270,5 @@ def draw_bounding_boxes(image, outputs):
 
     # Convert the OpenCV image (numpy array) to a PIL image
     image_pil = Image.fromarray(image_with_boxes_rgb)
-
+    image_pil.save(save_path, format="JPEG")
     return image_pil