Update app.py

app.py

@@ -68,5 +68,346 @@ model_cfg = "sam2_hiera_s.yaml"
 sam2_model = build_sam2(model_cfg, sam2_checkpoint, device="cuda")
 predictor = SAM2ImagePredictor(sam2_model)
 
-checkpoint_path = "/root/.cache/kagglehub/models/ybhavsar/newsegmentation/PyTorch/default/1/sam2_lr0.0001_wd0.01_900.torch"
-predictor.model.load_state_dict(torch.load(checkpoint_path))
+checkpoint_path = "sam2_lr0.0001_wd0.01_900.torch"
+predictor.model.load_state_dict(torch.load(checkpoint_path))
+
+def display_thread(thread_id):
+    for message in client.beta.threads.messages.list(thread_id=thread_id):
+        display(message.content[0].text.value)
+
+def read_file(filepath, max_pages=None):
+    if filepath.endswith('.pdf'):
+        return read_pdf(filepath, max_pages)
+    elif filepath.endswith('.txt'):
+        return read_text_file(filepath)
+    elif filepath.endswith('.docx'):
+        return read_docx(filepath)
+    elif filepath.endswith('.xlsx'):
+        return read_xlsx(filepath)
+    elif filepath.endswith('.pptx'):
+        return read_pptx(filepath)
+    else:
+        raise ValueError("Unsupported file type")
+
+def read_pdf(filepath, max_pages=None):
+    reader = PdfReader(filepath)
+    pdf_text = ""
+    page_number = 0
+    for page in reader.pages:
+        page_number += 1
+        if max_pages and (page_number > max_pages):
+            break
+        page_text = page.extract_text()
+
+        # Check if page_text is None before proceeding
+        if page_text:
+            # Replace multiple newlines with a space to make it readable
+            page_text = re.sub(r'\n+', ' ', page_text)
+            pdf_text += page_text + f"\nPage Number: {page_number}\n"
+        else:
+            pdf_text += f"\n[No extractable text on Page {page_number}]\n"
+
+    return pdf_text
+
+calc_similarity = lambda x, y: 1 - spatial.distance.cosine(x.data[0].embedding, y.data[0].embedding)
+
+def pretty_print(df):
+    return display(HTML(df.to_html().replace("\\n", "<br>")))
+
+
+def read_directory(directory):
+    assert os.path.exists(directory)
+
+    res_dict = {}
+
+    for filename in os.listdir(directory):
+        if filename.endswith(('pdf', 'txt', 'docx', 'pptx')):
+            filepath = os.path.join(directory, filename)
+            text = read_file(filepath, 2)
+            res_dict[filename] = (filepath, text)
+
+    df = pd.DataFrame(res_dict).T
+    df = df.reset_index()
+    df.columns = ["Filename", "Filepath", "Text"]
+    return df
+
+# Initialize GPT tokenizer
+tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
+tokenizer.model_max_length = int(1e30)
+
+def ask_chatbot(question, context, m):
+    max_context_tokens = 16385  # Adjust based on the maximum allowable context tokens
+    truncated_context = truncate_context(context, max_context_tokens)
+    response = client.chat.completions.create(
+        model=m,
+        messages=[
+            {"role": "system", "content": """You are an expert doctor who treats chronic wounds, and you know every single thing about wounds and how to treat them as well as preventing them from getting worse.
+    The user will provide the following inputs: Name, Gender, Age, Pre-existing Medical Conditions, Wound Part of Body, Wound Classficiation, Colors of the Wounds (as percents out of 100).
+    Please provide the medical advice in 2 concise paragraphs that must incorporate the following key features everytime:
+
+    1. **Wound Risk Score (1-100):** You will be given a PDF and you shall review it and use it to aid in your risk score generation. The wound risk score should be between 1-100! Of course, any color percentages **less than 3** shouldn't be taken into consideration when making the score.
+    **Make sure to be specific!**
+    2. **Medical Advice:** Give the patient bulleted directions on how to monitor and care for their wound. **Make sure to include if the person needs to go see a doctor as soon as possible.**"""},
+            {"role": "user", "content": truncated_context},
+            {"role": "user", "content": question}
+        ]
+    )
+    return response.choices[0].message.content
+
+def truncate_context(context, max_tokens):
+    tokens = tokenizer.encode(context)
+    if len(tokens) > max_tokens:
+        truncated_tokens = tokens[:max_tokens]
+        return tokenizer.decode(truncated_tokens)
+    return context
+
+
+file_content = read_file("Wound Healing Risk Assessment.pdf")
+
+api_key = "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" # @param {"type":"string"}
+
+%env OPENAI_API_KEY = {api_key}
+client = OpenAI()
+model="gpt-4o-mini"
+
+assistant = client.beta.assistants.create(
+    name="Wound Treater",
+    instructions="""You are an expert doctor who treats chronic wounds, and you know every single thing about wounds and how to treat them as well as preventing them from getting worse.
+    The user will provide the following inputs: Name, Gender, Age, Pre-existing Medical Conditions, Wound Part of Body, Wound Classficiation, Colors of the Wounds (as percents out of 100).
+    Please provide the medical advice in 2 concise paragraphs that must incorporate the following key features everytime:
+
+    1. **Wound Risk Score (1-100):** Generate a wound risk score from 1-100, 1 being no risk and 100 being going to see a medical professional immediately! Of course, any color percentages **less than 3** shouldn't be taken into consideration when making the score.
+    **Make sure to be specific and list the components of the wound risk score.**
+    2. **Medical Advice:** Give the patient directions on how to monitor and care for their wound. **Make sure to include if the person needs to go see a doctor as soon as possible.**""",
+    model=model)
+
+def get_assistant_response(name="None", gender="None", age="None", conditions="None", bodyPart="None", typeWound="None", red="None", orange="None", yellow="None", magenta="None", white="None", gray="None", black="None"):
+    thread = client.beta.threads.create()
+
+    input_text = "Name: " + str(name) + ", Gender: " + str(gender) + ", Age: " + str(age) + ", Pre-Existing Medical Conditions: " + str(conditions) + ", Part of Body: " + str(bodyPart) + ", Type of Wound: " + str(typeWound) + ", Wound Colors (Red, Orange, Yellow, Magenta, White, Gray, Black): [" + str(red) + ", " + str(orange) + ", " + str(yellow) + ", " + str(magenta) + ", " + str(white) + ", " + str(gray) + ", " + str(black) + "]"
+
+    message = client.beta.threads.messages.create(
+            thread_id=thread.id,
+            role="user",
+            content=input_text)
+    run = client.beta.threads.runs.create(
+          thread_id=thread.id,
+          assistant_id=assistant.id,
+        )
+    sleep(15)
+
+    return input_text, client.beta.threads.messages.list(thread.id).data[0].content[0].text.value
+
+def get_response_with_context(name="None", gender="None", age="None", conditions="None", bodyPart="None", typeWound="None", red="None", orange="None", yellow="None", magenta="None", white="None", gray="None", black="None"):
+    input_text = "Name: " + str(name) + ", Gender: " + str(gender) + ", Age: " + str(age) + ", Pre-Existing Medical Conditions: " + str(conditions) + ", Part of Body: " + str(bodyPart) + ", Type of Wound: " + str(typeWound) + ", Wound Colors (Red, Orange, Yellow, Magenta, White, Gray, Black): [" + str(red) + ", " + str(orange) + ", " + str(yellow) + ", " + str(magenta) + ", " + str(white) + ", " + str(gray) + ", " + str(black) + "]"
+
+    response = ask_chatbot(input_text, file_content, model)
+
+    return input_text, response
+
+
+wounds = []
+learn = load_learner('model.pkl')
+
+def one_step_inference(image_path, threshold=0.5):
+    image = cv2.imread(image_path)
+    image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+    image_hsv = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)
+
+    predictor.set_image(image)
+    high_res_features = [feat[-1].unsqueeze(0) for feat in predictor._features["high_res_feats"]]
+
+    with torch.no_grad():
+        sparse_embeddings, dense_embeddings = predictor.model.sam_prompt_encoder(points=None, boxes=None, masks=None)
+        low_res_masks, _, _, _ = predictor.model.sam_mask_decoder(
+            image_embeddings=predictor._features["image_embed"][-1].unsqueeze(0),
+            image_pe=predictor.model.sam_prompt_encoder.get_dense_pe(),
+            sparse_prompt_embeddings=sparse_embeddings,
+            dense_prompt_embeddings=dense_embeddings,
+            multimask_output=False,
+            repeat_image=False,  # Fixed argument
+            high_res_features=high_res_features,)
+
+    mask = predictor._transforms.postprocess_masks(low_res_masks, predictor._orig_hw[-1])
+    final_mask = (mask > threshold).cpu().detach().numpy()[0][0]
+    final_mask = final_mask.astype("uint8")
+
+    selected_pixels = cv2.bitwise_and(image_rgb, image_rgb, mask=final_mask)
+    selected_pixels = image_hsv[final_mask == 1]
+    colors = classify_colors(selected_pixels)
+
+    return colors["Red"], colors["Orange"], colors["Yellow"], colors["Magenta"], colors["White"], colors["Gray"], colors["Black"]
+
+# Replace this function in your original code
+def classify_colors(hsv_pixels):
+    # Define color ranges in HSV
+    color_ranges = {
+        'Red': [(0, 50, 50), (10, 255, 255)],  # Red wraps around
+        'Red2': [(170, 50, 50), (179, 255, 255)],
+        'Orange': [(11, 50, 50), (25, 255, 255)],
+        'Yellow': [(26, 50, 50), (35, 255, 255)],
+        'Green': [(36, 50, 50), (85, 255, 255)],
+        'Cyan': [(86, 50, 50), (95, 255, 255)],
+        'Blue': [(96, 50, 50), (130, 255, 255)],
+        'Purple': [(131, 50, 50), (160, 255, 255)],
+        'Magenta': [(161, 50, 50), (169, 255, 255)],
+        'White': [(0, 0, 200), (179, 55, 255)],  # High brightness, low saturation
+        'Gray': [(0, 0, 50), (179, 50, 200)],    # Low saturation, varying brightness
+        'Black': [(0, 0, 0), (179, 50, 50)]      # Low brightness
+    }
+    # Flatten the HSV pixels to process as a single list
+    hsv_pixels = hsv_pixels.reshape(-1, 3)
+    # Initialize counts for each color
+    color_counts = {color: 0 for color in color_ranges}
+    # Total number of pixels
+    total_pixels = hsv_pixels.shape[0]
+    # Classify each pixel
+    for pixel in hsv_pixels:
+        h, s, v = pixel
+        for color, ranges in color_ranges.items():
+            if isinstance(ranges[0], tuple):  # Handles multiple ranges (e.g., red)
+                lower = ranges[0]
+                upper = ranges[1]
+                if (lower[0] <= h <= upper[0] or lower[0] > upper[0] and (h >= lower[0] or h <= upper[0])) \
+                        and lower[1] <= s <= upper[1] and lower[2] <= v <= upper[2]:
+                    color_counts[color] += 1
+                    break
+            else:
+                lower, upper = ranges
+                if lower[0] <= h <= upper[0] and lower[1] <= s <= upper[1] and lower[2] <= v <= upper[2]:
+                    color_counts[color] += 1
+                    break
+    # Calculate percentages
+    color_counts["Red"] += color_counts["Red2"]
+    del color_counts["Red2"]
+    if(total_pixels == 0):
+        total_pixels = 1
+    color_percentages = {color: (count / total_pixels) * 100 for color, count in color_counts.items()}
+    return color_percentages
+
+def predict_image(image_path):
+    img = PILImage.create(image_path)
+    pred, pred_idx, probs = learn.predict(img)
+
+    return pred
+
+def reveal_group():
+    return gr.update(visible=True)
+
+def hide_group():
+    return gr.update(visible=False)
+
+# Add new wound to the list
+def add_wound(image, partOfBody):
+    wounds.append({"image": image, "description": partOfBody})
+    return image, partOfBody
+
+def clear_inputs(image, partOfBody):
+    image=None
+    partOfBody=""
+    return image, partOfBody
+
+# Initialize Gradio app
+with gr.Blocks(theme=gr.themes.Glass()) as demo:
+    gr.Markdown("<center><h1>Welcome to WoundView!</h1></center>")
+
+    # Sign-up Group
+    with gr.Group() as sign_up:
+        gr.Markdown("<center><h2>New User</h2></center>")
+        name = gr.Textbox(label="Full Name", placeholder="Enter your name here...")
+        gender = gr.Radio(["Male", "Female"], label="Gender")
+        age = gr.Number(label="Age")
+        conditions = gr.CheckboxGroup(["Diabetes", "Peripheral Arterial Disease", "Venous Insufficiency", "Obesity", "Smoking", ], label="Pre-Existing Medical Conditions")
+        gr.Markdown("<span style='color: red;'>Some fields were left empty. Please fill them out!</span>", visible=False)
+        sign_up_btn = gr.Button(value="Sign Up", variant="secondary")
+
+    # Home Group
+    with gr.Group(visible=False) as home:
+        gr.Markdown("<center><h2>Wounds</h2></center>")
+        with gr.Row(visible=False) as wound_display:
+            wound_image = gr.Image()
+            with gr.Column():
+                wound_title = gr.Markdown("<center><h2>Wound Description</h2></center>")
+                with gr.Row():
+                    gr.Markdown("<center>Part of Body:</center>")
+                    wound_desc = gr.Textbox(container=False)
+                with gr.Row():
+                    gr.Markdown("<center>Type of Wound:</center>")
+                    wound_classification = gr.Textbox(container=False)
+                gr.Markdown("<center><h4>Colors:</h4></center>")
+                with gr.Row():
+                    gr.Markdown("<center>Red:</center>")
+                    red_percent = gr.Textbox(container=False)
+                with gr.Row():
+                    gr.Markdown("<center>Orange:</center>")
+                    orange_percent = gr.Textbox(container=False)
+                with gr.Row():
+                    gr.Markdown("<center>Yellow:</center>")
+                    yellow_percent = gr.Textbox(container=False)
+                with gr.Row():
+                    gr.Markdown("<center>Magenta:</center>")
+                    magenta_percent = gr.Textbox(container=False)
+                with gr.Row():
+                    gr.Markdown("<center>White:</center>")
+                    white_percent = gr.Textbox(container=False)
+                with gr.Row():
+                    gr.Markdown("<center>Gray:</center>")
+                    gray_percent = gr.Textbox(container=False)
+                with gr.Row():
+                    gr.Markdown("<center>Black:</center>")
+                    black_percent = gr.Textbox(container=False)
+                ai_chat_btn = gr.Button(value="AI ChatBot")
+        add_new_btn = gr.Button(value="Add New")
+
+    # Add New Group
+    with gr.Group(visible=False) as add_new:
+        gr.Markdown("<center><h2>Add New Wound</h2></center>")
+        with gr.Row():
+            image = gr.Image(label="Picture of wound", type="filepath")
+            partOfBody = gr.Radio(["Head", "Arm", "Hand", "Back", "Stomach", "Leg", "Foot"], label="What part of the body is the wound on?")
+        with gr.Row():
+            confirm_add_new_btn = gr.Button(value="Confirm")
+            cancel_add_new_btn = gr.Button(value="Cancel")
+
+    with gr.Group(visible=False) as ai_chat:
+        gr.Markdown("<center><h2>AI Chat</h2></center>")
+        with gr.Column() as gpt:
+            gr.Markdown("<center><h3>Chat GPT</h3></center>")
+            chatGPTInput = gr.Textbox(container=False)
+            chatGPTOutput = gr.Textbox(container=False)
+        cancel_ai_chat_btn = gr.Button(value="Cancel")
+
+    # Button Click Events
+    sign_up_btn.click(hide_group, outputs=sign_up).then(reveal_group, outputs=home)
+
+    add_new_btn.click(hide_group, outputs=home).then(reveal_group, outputs=add_new
+    ).then(clear_inputs,
+        inputs=[image, partOfBody],
+        outputs=[image, partOfBody]
+    )
+
+    confirm_add_new_btn.click(add_wound,
+        inputs=[image, partOfBody],
+        outputs=[wound_image, wound_desc]
+    ).then(reveal_group, outputs=home
+    ).then(hide_group, outputs=add_new
+    ).then(reveal_group, outputs=wound_display
+    ).then(predict_image,
+        inputs=image,
+        outputs=wound_classification
+    ).then(one_step_inference,
+        inputs=image,
+        outputs=[red_percent, orange_percent, yellow_percent, magenta_percent, white_percent, gray_percent, black_percent]
+    )
+
+    cancel_add_new_btn.click(hide_group, outputs=add_new).then(reveal_group, outputs=home)
+
+    ai_chat_btn.click(hide_group, outputs=home).then(reveal_group, outputs=ai_chat
+    ).then(get_response_with_context,
+        inputs=[name, gender, age, conditions, partOfBody, wound_classification, red_percent, orange_percent, yellow_percent, magenta_percent, white_percent, gray_percent, black_percent],
+        outputs=[chatGPTInput, chatGPTOutput]
+    )
+
+    cancel_ai_chat_btn.click(hide_group, outputs=ai_chat).then(reveal_group, outputs=home)
+
+demo.launch(share=True)