BAAI/SegVol

1import matplotlib.pyplot as plt2import numpy as np3from PIL import Image, ImageEnhance, ImageDraw4import torch5import streamlit as st6from model.inference_cpu import inference_case7 8initial_rectangle = {9    "version": "4.4.0",10    'objects': [11        {12            "type": "rect",13            "version": "4.4.0",14            "originX": "left",15            "originY": "top",16            "left": 50,17            "top": 50,18            "width": 100,19            "height": 100,20            'fill': 'rgba(255, 165, 0, 0.3)', 21            'stroke': '#2909F1', 22            'strokeWidth': 3, 23            'strokeDashArray': None, 24            'strokeLineCap': 'butt', 25            'strokeDashOffset': 0, 26            'strokeLineJoin': 'miter', 27            'strokeUniform': True, 28            'strokeMiterLimit': 4, 29            'scaleX': 1, 30            'scaleY': 1, 31            'angle': 0, 32            'flipX': False, 33            'flipY': False, 34            'opacity': 1, 35            'shadow': None, 36            'visible': True, 37            'backgroundColor': '', 38            'fillRule': 39            'nonzero', 40            'paintFirst': 41            'fill', 42            'globalCompositeOperation': 'source-over', 43            'skewX': 0, 44            'skewY': 0, 45            'rx': 0, 46            'ry': 047        }48    ]49}50 51def run():52    image = st.session_state.data_item["image"].float()53    image_zoom_out = st.session_state.data_item["zoom_out_image"].float()54    text_prompt = None55    point_prompt = None56    box_prompt = None57    if st.session_state.use_text_prompt:58        text_prompt = st.session_state.text_prompt59    if st.session_state.use_point_prompt and len(st.session_state.points) > 0:60        point_prompt = reflect_points_into_model(st.session_state.points)61    if st.session_state.use_box_prompt:62        box_prompt = reflect_box_into_model(st.session_state.rectangle_3Dbox)63    inference_case.clear()64    st.session_state.preds_3D, st.session_state.preds_3D_ori = inference_case(image, image_zoom_out, 65                                            text_prompt=text_prompt,66                                            _point_prompt=point_prompt,67                                            _box_prompt=box_prompt)68 69def reflect_box_into_model(box_3d):70    z1, y1, x1, z2, y2, x2 = box_3d71    x1_prompt = int(x1 * 256.0 / 325.0)72    y1_prompt = int(y1 * 256.0 / 325.0)73    z1_prompt = int(z1 * 32.0 / 325.0)74    x2_prompt = int(x2 * 256.0 / 325.0)75    y2_prompt = int(y2 * 256.0 / 325.0)76    z2_prompt = int(z2 * 32.0 / 325.0)77    return torch.tensor(np.array([z1_prompt, y1_prompt, x1_prompt, z2_prompt, y2_prompt, x2_prompt]))78 79def reflect_json_data_to_3D_box(json_data, view):80    if view == 'xy':81        st.session_state.rectangle_3Dbox[1] = json_data['objects'][0]['top']82        st.session_state.rectangle_3Dbox[2] = json_data['objects'][0]['left']83        st.session_state.rectangle_3Dbox[4] = json_data['objects'][0]['top'] + json_data['objects'][0]['height'] * json_data['objects'][0]['scaleY']84        st.session_state.rectangle_3Dbox[5] = json_data['objects'][0]['left'] + json_data['objects'][0]['width'] * json_data['objects'][0]['scaleX']85    print(st.session_state.rectangle_3Dbox)86 87def reflect_points_into_model(points):88    points_prompt_list = []89    for point in points:90        z, y, x = point91        x_prompt = int(x * 256.0 / 325.0)92        y_prompt = int(y * 256.0 / 325.0)93        z_prompt = int(z * 32.0 / 325.0)94        points_prompt_list.append([z_prompt, y_prompt, x_prompt])95    points_prompt = np.array(points_prompt_list)96    points_label = np.ones(points_prompt.shape[0])97    print(points_prompt, points_label)98    return (torch.tensor(points_prompt), torch.tensor(points_label))99 100def show_points(points_ax, points_label, ax):101    color = 'red' if points_label == 0 else 'blue'102    ax.scatter(points_ax[0], points_ax[1], c=color, marker='o', s=200)103 104def make_fig(image, preds, point_axs=None, current_idx=None, view=None):105    # Convert A to an image106    image = Image.fromarray((image * 255).astype(np.uint8)).convert("RGB")107    enhancer = ImageEnhance.Contrast(image)108    image = enhancer.enhance(2.0)109 110    # Create a yellow mask from B111    if preds is not None:112        mask = np.where(preds == 1, 255, 0).astype(np.uint8)113        mask = Image.merge("RGB", 114                           (Image.fromarray(mask), 115                            Image.fromarray(mask), 116                            Image.fromarray(np.zeros_like(mask, dtype=np.uint8))))117 118        # Overlay the mask on the image119        image = Image.blend(image.convert("RGB"), mask, alpha=st.session_state.transparency)120    121    if point_axs is not None:122        draw = ImageDraw.Draw(image)123        radius = 5124        for point in point_axs:125            z, y, x = point126            if view == 'xy' and z == current_idx:127                draw.ellipse((x-radius, y-radius, x+radius, y+radius), fill="blue")128            elif view == 'xz'and y == current_idx:129                draw.ellipse((x-radius, z-radius, x+radius, z+radius), fill="blue")130    return image