--- a +++ b/webapp.py @@ -0,0 +1,134 @@ +import streamlit as st +import cv2 +import numpy as np +import tempfile +import os +from ultralytics import YOLO + +# Load the YOLOv8 model +try: + model = YOLO('best.pt') # Ensure the path to your trained YOLOv8 model weights is correct +except FileNotFoundError: + st.error("Model weights file 'best.pt' not found. Please ensure the model file is in the correct directory.") + st.stop() + +# Predefined colors for each class +COLORS = { + "high-pneumonia": [255, 0, 0], # Red + "low-pneumonia": [0, 255, 0], # Green + "no-pneumonia": [0, 0, 255] # Blue +} + +# Define a constant image size +IMAGE_SIZE = (640, 640) # Width, Height + +def segment_image(image): + results = model(image) + return results + +def display_segmented_image(image, results): + detected_classes = set() + if results and results[0].masks is not None and results[0].boxes is not None: + masks = results[0].masks.data.cpu().numpy() + boxes = results[0].boxes.data.cpu().numpy() + class_ids = results[0].boxes.cls.cpu().numpy() + + for mask, box, class_id in zip(masks, boxes, class_ids): + class_name = model.names[int(class_id)] + detected_classes.add(class_name) + + # If low-pneumonia or high-pneumonia is detected, ignore no-pneumonia masks + if 'low-pneumonia' in detected_classes or 'high-pneumonia' in detected_classes: + detected_classes.discard('no-pneumonia') + + for mask, box, class_id in zip(masks, boxes, class_ids): + class_name = model.names[int(class_id)] + if class_name not in detected_classes: + continue + mask = mask.astype(np.uint8) + mask_resized = cv2.resize(mask, (image.shape[1], image.shape[0]), interpolation=cv2.INTER_NEAREST) + color = COLORS.get(class_name, [255, 255, 255]) # Default to white if class not found + image[mask_resized == 1] = color + x1, y1, x2, y2 = box[:4].astype(int) + # Draw a rectangle behind the text for better visibility + text_size = cv2.getTextSize(class_name, cv2.FONT_HERSHEY_SIMPLEX, 0.9, 2)[0] + text_x = x1 + text_y = y1 - 10 + cv2.rectangle(image, (text_x, text_y - text_size[1]), (text_x + text_size[0], text_y), color, -1) + cv2.putText(image, class_name, (text_x, text_y), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (255, 255, 255), 2) + # Draw the bounding box rectangle + cv2.rectangle(image, (x1, y1), (x2, y2), color, 2) + else: + st.warning("No objects detected.") + + # Resize the final image to the constant size + resized_image = cv2.resize(image, IMAGE_SIZE, interpolation=cv2.INTER_LINEAR) + + return resized_image, detected_classes + + +st.title('X-ray Segmentation Project') + +# Add an author section in the sidebar +st.sidebar.title('About the Author') +st.sidebar.markdown(""" + **Author Name**: Makhammadjonov Izzatullokh + **Email**: izzatullokhm@gmail.com + <a href="https://github.com/Izzatullokh24" target="_blank"><i class="fab fa-github"></i> GitHub</a> + <a href="https://www.linkedin.com/in/izzatullokh-makhammadjonov-242042195/" target="_blank"><i class="fab fa-linkedin"></i> LinkedIn</a> + <style> + .fab { + font-size: 24px; + margin-right: 10px; + } + </style> + <link rel="stylesheet" href="https://cdnjs.cloudflare.com/ajax/libs/font-awesome/5.15.4/css/all.min.css"> + Izzatullokh is a machine learning engineer with a passion for computer vision and deep learning. +""", unsafe_allow_html=True) + +uploaded_file = st.file_uploader("Upload an image or video", type=["jpg", "jpeg", "png", "mp4", "avi"]) + +if uploaded_file is not None: + file_type = uploaded_file.type.split('/')[0] + + try: + if file_type == 'image': + file_bytes = np.asarray(bytearray(uploaded_file.read()), dtype=np.uint8) + image = cv2.imdecode(file_bytes, cv2.IMREAD_COLOR) + results = segment_image(image) + segmented_image, detected_classes = display_segmented_image(image, results) + st.image(segmented_image, caption='Segmented Image', use_column_width=True) + + # Display pneumonia information + if detected_classes: + st.subheader('Diagnosis:') + diagnosis_html = "" + if 'high-pneumonia' in detected_classes: + diagnosis_html += '<p style="color: red; font-size: 20px;">The patient has <strong>high pneumonia</strong>.</p>' + if 'low-pneumonia' in detected_classes: + diagnosis_html += '<p style="color: green; font-size: 20px;">The patient has <strong>low pneumonia</strong>.</p>' + if 'no-pneumonia' in detected_classes: + diagnosis_html += '<p style="color: blue; font-size: 20px;">The patient does <strong>not have pneumonia</strong>.</p>' + st.markdown(diagnosis_html, unsafe_allow_html=True) + else: + st.write('No pneumonia detected.') + + elif file_type == 'video': + tfile = tempfile.NamedTemporaryFile(delete=False) + tfile.write(uploaded_file.read()) + video_path = tfile.name + segmented_frames = segment_video(video_path) + stframe = st.empty() + for frame, results in segmented_frames: + segmented_frame, detected_classes = display_segmented_image(frame, results) + stframe.image(segmented_frame, channels="BGR") + os.remove(video_path) + + else: + st.error("Unsupported file format. Please upload a jpg, jpeg, png, mp4, or avi file.") + + except Exception as e: + st.error(f"An error occurred: {e}") + +else: + st.info("Please upload a file to proceed.")