first commit

app.py

@@ -0,0 +1,26 @@
+from heinsight import HeinSight, HeinSightConfig
+
+heinsight = HeinSight(vial_model_path="models/best_vessel.pt",
+                      contents_model_path="models/best_content.pt",
+                      config=HeinSightConfig())
+
+import gradio as gr
+
+# Gradio UI
+demo = gr.Interface(
+    fn=heinsight.image_demo,
+    inputs=[
+        gr.Image(type="pil"),
+        gr.Slider(0.1, 1.0, step=0.01, value=0, label="Cap Size Ratio")
+    ],
+    outputs=[
+        gr.Image(type="pil", label="Detected Image"),
+        gr.JSON(label="Detection Info")  # or gr.Textbox() if you prefer plain text
+    ],
+    title="HeinSight",
+    description="Upload an image with vials to detect their contents"
+)
+
+
+if __name__ == "__main__":
+    demo.launch()

heinsight.py

@@ -0,0 +1,303 @@
+
+from itertools import chain
+from random import randint
+
+import cv2
+import matplotlib
+import matplotlib.pyplot as plt
+import numpy as np
+import pandas as pd
+from ultralytics import YOLO
+from PIL import Image
+matplotlib.use('Agg')
+
+def highlight_vial_body(frame, vial_location, cap_ratio=0.2):
+    """
+    Highlights only the vial body in the frame by masking out background and cap.
+
+    Args:
+        frame (np.ndarray): Original BGR image.
+        vial_location (tuple): (x, y, w, h) bounding box of vial.
+        cap_ratio (float): Fraction (0-1) of vial height considered as cap.
+
+    Returns:
+        masked_frame (np.ndarray): Frame with background and cap masked out.
+    """
+    overlay = frame.copy()
+    x, y, x2, y2 = vial_location
+    h = y2 - y
+    # Define cap and body regions
+    cap_height = int(h * cap_ratio)
+    body_y_start = y + cap_height
+
+    # Draw gray background mask
+    cv2.rectangle(overlay, (0, 0), (frame.shape[1], frame.shape[0]), (128, 128, 128), thickness=-1)
+
+    # Draw red translucent cap over the vial's cap region
+    cv2.rectangle(overlay, (x, y), (x2, body_y_start), (0, 0, 255), thickness=-1)
+    masked = cv2.addWeighted(overlay, 0.5, frame, 0.5, 0)
+
+    return masked
+
+
+class HeinSightConfig:
+
+    """Configuration for the HeinSight system."""
+    NUM_ROWS = -1
+    SAVE_PLOT_VIDEO = True
+    LIQUID_CONTENT = ["Homo", "Hetero"]
+    CAP_RATIO = 0.3
+    STATUS_RULE = 0.7
+    DEFAULT_VIAL_LOCATION = None
+    DEFAULT_VIAL_HEIGHT = None
+
+class HeinSight:
+    """
+    The core of the HeinSight system, responsible for computer vision and analysis.
+    """
+
+    def __init__(self, vial_model_path: str, contents_model_path: str, config: HeinSightConfig = HeinSightConfig()):
+        self.fig, self.axs = plt.subplots(2, 2, figsize=(8, 6), height_ratios=[2, 1], constrained_layout=True)
+        self._set_axes()
+        self.config = config
+        self.vial_model = YOLO(vial_model_path)
+        self.contents_model = YOLO(contents_model_path)
+        self.color_palette = self._register_colors([self.vial_model, self.contents_model])
+        self.clear_cache()
+
+    def _set_axes(self):
+        """creating plot axes"""
+        ax0, ax1, ax2, ax3 = self.axs.flat
+        ax0.set_position([0.21, 0.45, 0.22, 0.43])  # [left, bottom, width, height]
+
+        ax1.set_position([0.47, 0.45, 0.45, 0.43])  # [left, bottom, width, height]
+        ax2.set_position([0.12, 0.12, 0.35, 0.27])
+        ax3.set_position([0.56, 0.12, 0.35, 0.27])
+        self.fig.canvas.draw_idle()
+
+    def clear_cache(self):
+        """Resets the state of the HeinSight system."""
+        self.vial_location = self.config.DEFAULT_VIAL_LOCATION.copy() if self.config.DEFAULT_VIAL_LOCATION else None
+        self.cap_rows = 0
+        self.vial_heigh = self.config.DEFAULT_VIAL_HEIGHT
+        self.vial_size = []
+        self.content_info = None
+        self.x_time = []
+        self.turbidity_2d = []
+        self.average_colors = []
+        self.average_turbidity = []
+        self.output = []
+        self.stream_output = []
+        self.status = {}
+        self.status_queue = []
+        self.output_dataframe = pd.DataFrame()
+        self.output_frame = None
+        self.turbidity = []
+
+    @staticmethod
+    def _register_colors(model_list):
+        """
+        register default colors for models
+        :param model_list: YOLO models list
+        """
+        name_color_dict = {
+            "Empty": (19, 69, 139),  # Brown
+            "Residue": (0, 165, 255),  # Orange
+            "Hetero": (255, 0, 255),  # purple
+            "Homo": (0, 0, 255),  # Red
+            "Solid": (255, 0, 0),  # Blue
+        }
+        names = set(chain.from_iterable(model.names.values() for model in model_list if model))
+        for name in names:
+            if name not in name_color_dict:
+                name_color_dict[name] = (randint(0, 255), randint(0, 255), randint(0, 255))
+        return name_color_dict
+
+    def find_vial(self, frame):
+        """
+        Detect the vial in video frame with YOLOv8
+        :param frame: raw input frame
+        :return result: np.ndarray or None: Detected vial bounding box or None if no vial is found.
+        """
+        # vial location is not defined, use vial model to detect
+        if not self.vial_location:
+            results = self.vial_model(frame, conf=0.2, max_det=1)
+            boxes = results[0].boxes.data.cpu().numpy()
+            if boxes.size > 0:
+                self.vial_location = [int(x) for x in boxes[0, :4]]
+        if self.vial_location:
+            self.cap_rows = int((self.vial_location[3] - self.vial_location[1]) * self.config.CAP_RATIO)
+        return self.vial_location is not None
+
+    def crop_rectangle(self, image, vial_location):
+        """
+        crop and resize the image
+        :param image: raw image capture
+        :param vial_location:
+        :return: cropped and resized vial frame
+        """
+        x1, y1, x2, y2 = vial_location
+        y1 = int(self.config.CAP_RATIO * (y2 - y1)) + y1
+        cropped_image = image[y1:y2, x1:x2]
+        return cropped_image
+
+    def content_detection(self, vial_frame):
+        """
+        Detect content in a vial frame.
+        :param vial_frame: (np.ndarray) Cropped vial frame.
+        :return tuple: Bounding boxes, liquid boxes, and detected class titles.
+        """
+        results = self.contents_model(vial_frame, max_det=4, agnostic_nms=False, conf=0.25, iou=0.25, verbose=False)
+        bboxes = results[0].boxes.data.cpu().numpy()
+        pred_classes = bboxes[:, 5]
+        title = " ".join([self.contents_model.names[int(x)] for x in pred_classes])
+        liquid_boxes = [bboxes[i][:4] for i, cls in enumerate(pred_classes) if
+                        self.contents_model.names[int(cls)] in self.config.LIQUID_CONTENT]
+        return bboxes, sorted(liquid_boxes, key=lambda x: x[1], reverse=True), title
+
+
+    def process_vial_frame(self, vial_frame, update_od: bool = False):
+        """
+        process single vial frame, detect content, draw bounding box and calculate turbidity and color
+        :param vial_frame: vial frame image
+        :param update_od: update object detection, True: run YOLO for this frame, False: use previous YOLO results
+        """
+        if update_od or self.content_info is None:
+            self.content_info = self.content_detection(vial_frame)
+        bboxes, liquid_boxes, title = self.content_info
+        phase_data, raw_turbidity = self.calculate_value_color(vial_frame, liquid_boxes)
+        frame_image = self.draw_bounding_boxes(vial_frame, bboxes, self.contents_model.names, text_right=False)
+
+        if self.config.SAVE_PLOT_VIDEO:
+            self.display_frame(raw_turbidity, frame_image, title)
+            self.fig.canvas.draw()
+            frame_image = np.array(self.fig.canvas.renderer.buffer_rgba())
+            frame_image = cv2.cvtColor(frame_image, cv2.COLOR_RGBA2BGR)
+        return frame_image, bboxes, raw_turbidity, phase_data
+
+    def calculate_value_color(self, vial_frame, liquid_boxes):
+        """
+        Calculate the value and color for a given vial image and bounding boxes
+        :param vial_frame: the vial image
+        :param liquid_boxes: the liquid boxes (["Homo", "Hetero"])
+        :return: the output dict and raw turbidity per row
+        """
+        height, _, _ = vial_frame.shape
+        hsv_image = cv2.cvtColor(vial_frame, cv2.COLOR_BGR2HSV)
+        output = {
+            'time': self.x_time[-1],
+            'color': np.mean(hsv_image[:, :, 0]),
+            'turbidity': np.mean(hsv_image[:, :, 2])
+        }
+        raw_value = np.mean(hsv_image[:, :, 2], axis=1)
+        for i, bbox in enumerate(liquid_boxes):
+            _, top, _, bottom = map(int, bbox)
+            roi = hsv_image[top:bottom, :]
+            output[f'volume_{i + 1}'] = (bottom - top) / height
+            output[f'color_{i + 1}'] = np.mean(roi[:, :, 0])
+            output[f'turbidity_{i + 1}'] = np.mean(roi[:, :, 2])
+        self.average_colors.append(output['color'])
+        self.average_turbidity.append(output['turbidity'])
+        return output, raw_value
+
+    @staticmethod
+    def _get_dynamic_font_params(img_height, base_height=200, base_font_scale=0.5, base_thickness=1):
+        scale_factor = img_height / base_height
+        font_scale = base_font_scale * scale_factor
+        thickness = max(1, int(base_thickness * scale_factor))
+        return font_scale, thickness
+
+    def draw_bounding_boxes(self, image, bboxes, class_names, thickness=None, text_right=False, on_raw=False):
+        """Draws bounding boxes on the image."""
+        output_image = image.copy()
+        height = image.shape[1]
+        font_scale, text_thickness = self._get_dynamic_font_params(height)
+        margin = 2
+        thickness = thickness or max(1, int(height / 200))
+        for rect in bboxes:
+            x1, y1, x2, y2, _, class_id = map(int, rect)
+            class_name = class_names[class_id]
+            color = self.color_palette.get(class_name, (255, 255, 255))
+            if on_raw and self.vial_location:
+                x1, y1 = x1 + self.vial_location[0], y1 + self.vial_location[1] + self.cap_rows
+                x2, y2 = x2 + self.vial_location[0], y2 + self.vial_location[1] + self.cap_rows
+            cv2.rectangle(output_image, (x1, y1), (x2, y2), color, thickness)
+            (text_width, text_height), baseline = cv2.getTextSize(class_name, cv2.FONT_HERSHEY_SIMPLEX, font_scale,
+                                                                  text_thickness)
+            text_location = (
+                x2 - text_width - margin if text_right ^ (class_name == "Solid") else x1 + margin,
+                y1 + text_height + margin
+            )
+            cv2.putText(output_image, class_name, text_location, cv2.FONT_HERSHEY_SIMPLEX, font_scale, color,
+                        text_thickness)
+        return output_image
+
+    def display_frame(self, y_values, image, title=None):
+        """
+        Display the image (top-left) and its turbidity values per row (top-right)
+        turbidity over time (bottom-left) and color over time (bottom-right)
+        :param y_values: the turbidity value per row
+        :param image: vial image frame to display
+        :param title: title of the image frame
+        """
+        # init plot
+        for ax in self.axs.flat:
+            ax.clear()
+        ax0, ax1, ax2, ax3 = self.axs.flat
+
+        # top left - vial frame and bounding boxes
+        image_copy = image.copy()
+        image_copy = cv2.cvtColor(image_copy, cv2.COLOR_BGR2RGB)
+        ax0.imshow(np.flipud(image_copy), origin='lower')
+        if title:
+            ax0.set_title(title)
+
+        # use fill between to optimize the speed 154.9857677 -> 68.15193
+        x_values = np.arange(len(y_values))
+        ax1.fill_betweenx(x_values, 0, y_values[::-1], color='green', alpha=0.5)
+        ax1.set_ylim(0, len(y_values))
+        ax1.set_xlim(0, 255)
+        ax1.xaxis.set_label_position('top')
+        ax1.set_xlabel('Turbidity per row')
+
+        realtime_tick_label = None
+
+        # bottom left - turbidity
+        ax2.set_ylabel('Turbidity')
+        ax2.set_xlabel('Time / min')
+        ax2.plot(self.x_time, self.average_turbidity)
+        ax2.set_xticks([self.x_time[0], self.x_time[-1]], realtime_tick_label)
+
+
+
+        # bottom right - color
+        ax3.set_ylabel('Color (hue)')
+        ax3.set_xlabel('Time / min')
+        ax3.plot(self.x_time, self.average_colors)
+        ax3.set_xticks([self.x_time[0], self.x_time[-1]], realtime_tick_label)
+
+
+    def image_demo(self, pil_image, cap_ratio=0):
+
+        self.clear_cache()
+        frame = cv2.cvtColor(np.array(pil_image), cv2.COLOR_RGB2BGR)  # PIL → OpenCV
+        phase_data = {}
+        self.config.CAP_RATIO = cap_ratio
+        if self.find_vial(frame):
+            vial_frame = self.crop_rectangle(frame, self.vial_location)
+            x1, y1, x2, y2 = self.vial_location
+            self.x_time.append(0)
+            frame_image, bboxes, _, phase_data = self.process_vial_frame(vial_frame)
+            boxes_on_vial = self.draw_bounding_boxes(vial_frame, bboxes, self.contents_model.names, on_raw=False)
+            masked_frame = highlight_vial_body(frame, self.vial_location, cap_ratio=cap_ratio)
+            masked_frame[y1 + self.cap_rows :y2, x1:x2] = boxes_on_vial
+
+            # bboxes_on_raw = self.draw_bounding_boxes(masked_frame, bboxes, self.contents_model.names, on_raw=True)
+            result = masked_frame
+        else:
+            result = frame
+        result_rgb = cv2.cvtColor(result, cv2.COLOR_BGR2RGB)  # OpenCV → RGB
+        return Image.fromarray(result_rgb), phase_data
+
+
+

models/best_content.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:00b0be2dd8eec4aedd5d56da0fa196b95454d420a0df7eae453178ab9fbcc485
+size 52009878

models/best_vessel.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:93a49e5f80434b35f9244ab67fb17f5440be84ec0427834dff28ad98fa83bc58
+size 22496110

requirements.txt

@@ -0,0 +1,3 @@
+Ultralytics
+Pillow
+gradio