import os import base64 import json import streamlit as st from streamlit_option_menu import option_menu from streamlit_authenticator import Authenticate import yaml from yaml.loader import SafeLoader import pandas as pd from PIL import Image import numpy as np import torch import cv2 from transformers import AutoImageProcessor, AutoModelForObjectDetection import bbox_visualizer as bbv from st_clickable_images import clickable_images from glob import glob MODEL_PATH = "pankaj-munde/FScout_v0.2" # image_dir = "./Data/images/" detr_preprocessor = AutoImageProcessor.from_pretrained(MODEL_PATH, token=st.secrets["HF_TOKEN"]) detr_model = AutoModelForObjectDetection.from_pretrained(MODEL_PATH, token=st.secrets["HF_TOKEN"]) colors = [[236, 112, 99], [165, 105, 189], [ 225, 9, 232], [ 255, 38, 8 ], [ 247, 249, 249 ], [170, 183, 184 ], [ 247, 249, 249 ], [ 247, 249, 249 ]] # with open('./static/config.yaml') as file: # config = yaml.load(file, Loader=SafeLoader) config = json.loads(st.secrets["CONFIG"]) authenticator = Authenticate( config['credentials'], config['cookie']['name'], config['cookie']['key'], config['cookie']['expiry_days'], config['preauthorized'] ) name, authentication_status, username = authenticator.login('Login', 'main') # images_lst = os.listdir(image_dir) # images = [] # for file in images_lst: # ipath = os.path.join(os.path.abspath(image_dir), file) # with open(ipath, "rb") as image: # encoded = base64.b64encode(image.read()).decode() # images.append(f"data:image/jpeg;base64,{encoded}") def get_detr_predictions(image, thresh): with torch.no_grad(): inputs = detr_preprocessor(images=image, return_tensors="pt") outputs = detr_model(**inputs) target_sizes = torch.tensor([image.size[::-1]]) results = detr_preprocessor.post_process_object_detection( outputs, threshold=float(thresh), target_sizes=target_sizes)[0] return results def add_label(img, label, bbox, draw_bg=True, text_bg_color=(255, 255, 255), text_color=(0, 0, 0), top=True): """adds label, inside or outside the rectangle Parameters ---------- img : ndarray the image on which the label is to be written, preferably the image with the rectangular bounding box drawn label : str the text (label) to be written bbox : list a list containing x_min, y_min, x_max and y_max of the rectangle positions draw_bg : bool, optional if True, draws the background of the text, else just the text is written, by default True text_bg_color : tuple, optional the background color of the label that is filled, by default (255, 255, 255) text_color : tuple, optional color of the text (label) to be written, by default (0, 0, 0) top : bool, optional if True, writes the label on top of the bounding box, else inside, by default True Returns ------- ndarray the image with the label written """ text_width = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 1, 2)[0][0] if top: label_bg = [bbox[0], bbox[1], bbox[0] + text_width, bbox[1] + 30] if draw_bg: cv2.rectangle(img, (label_bg[0], label_bg[1]), (label_bg[2] + 5, label_bg[3]), text_bg_color, -1) cv2.putText(img, label, (bbox[0] + 5, bbox[1] - 5), cv2.FONT_HERSHEY_SIMPLEX, 1, text_color, 2) else: label_bg = [bbox[0], bbox[1], bbox[0] + text_width, bbox[1] + 30] if draw_bg: cv2.rectangle(img, (label_bg[0], label_bg[1]), (label_bg[2] + 5, label_bg[3]), text_bg_color, -1) cv2.putText(img, label, (bbox[0] + 5, bbox[1] - 5 + 30), cv2.FONT_HERSHEY_SIMPLEX, 1, text_color, 2) return img def image_checkup_v4(ipath, thresh, show_count): imgOrig = ipath.convert("RGB") image = imgOrig.copy() # crop_name, crop_conf, crop_id = get_crop(image) detr_results = get_detr_predictions(image, thresh) final_results = [] all_predictions = { "Name": "Detailed View", "Value": "" } # result_data = {crop_name: crop_conf, "Inspection_data": []} img_with_box = np.array(image).copy() for idx, label_id in enumerate(detr_results["labels"].numpy()): pred_score = round(detr_results["scores"].numpy()[idx], 2) predicted_label = detr_model.config.id2label[label_id] # if float(pred_score) > 50: bbox = list(np.array(detr_results["boxes"].numpy()[idx], dtype=int)) img_with_box = bbv.draw_rectangle(img_with_box, bbox, bbox_color=colors[label_id]) # img_with_box = bbv.add_label(img_with_box, label=f"{predicted_label} : {pred_score}", bbox=bbox, top=False) if show_count: img_with_box = bbv.add_label( img_with_box, f"{idx + 1}", bbox, draw_bg=True, top=True) else: img_with_box = bbv.add_label( img_with_box, f"", bbox, draw_bg=False, top=True) final_results.append( {"prediction": predicted_label, "confidence": pred_score, "color": colors[label_id] } ) all_predictions["Value"] += f"\n{idx + 1}. {predicted_label.split('_')[-1]} - {round(pred_score, 2)}%\n" if len(final_results) > 0: df = pd.DataFrame(final_results) info = df["prediction"].value_counts() resized_seg = cv2.resize(img_with_box, imgOrig.size) new_res = [] for k, v in dict(info).items(): tmp = {} prd_id = detr_model.config.label2id[k] tmp["Insect"] = k tmp["Count"] = v tmp["Color"] = colors[int(prd_id)] new_res.append(tmp) return new_res, resized_seg return [], img_with_box def add_logo(logo_path, width, height): """Read and return a resized logo""" logo = Image.open(logo_path) # modified_logo = logo.resize((width, height)) return logo # st.write("
", unsafe_allow_html=True) if st.session_state["authentication_status"]: with st.sidebar: my_logo = add_logo(logo_path="./static//FarmGyan logo_1.png", width=50, height=60) st.image(my_logo) ucol, bcol = st.columns([3, 2]) ucol.write(f'## Welcome *{st.session_state["name"]}*') with bcol: authenticator.logout('Logout', 'main') st.write("
", unsafe_allow_html=True) st.title(":seedling: FarmGyan | Insects Scouting") st.write("
", unsafe_allow_html=True) st.write("## 🖖 Upload image for prediction") uploaded_file = st.file_uploader("Choose an image file", type=["jpg", "jpeg", "png"]) st.write("
", unsafe_allow_html=True) with st.spinner(text='In progress'): st.sidebar.write("## ⚙️ Configurations") st.sidebar.write("
", unsafe_allow_html=True) st.sidebar.write("#### Prediction Threshold") thresh = st.sidebar.slider("Threshold", 0.0, 1.0, 0.7, 0.1) st.sidebar.write("#### Boxes Count") show_count = st.sidebar.checkbox("Show Count") if uploaded_file is not None: clicked = None image = Image.open(uploaded_file).convert("RGB") predicted_data, result_image = image_checkup_v4(image, thresh, show_count) # print(predicted_data) col, col1 = st.columns([2, 4]) feedback_submitted = False # Initialize the flag with col: st.subheader("💯 Predicted Labels") st.write(f"

Total Count : {sum([d['Count'] for d in predicted_data])}

", unsafe_allow_html=True) for i, d in enumerate(predicted_data): # Create HTML markup with style information html_string = f"""
{i + 1}.

{d["Insect"]} : {d["Count"]}

""" st.markdown(html_string, unsafe_allow_html=True) st.write("
", unsafe_allow_html=True) with col1: st.subheader("🍀 Predicted Image") st.write("
", unsafe_allow_html=True) st.image(result_image)